cubecl-tile 0.12.0-pre.1

//! A paged KV cache, which is expert routing with the axis split.
//!
//! The physical pool is `[PAGE, OFFSET]` over one buffer, spelled with
//! [`PhysicalAxisMap::disjoint`] exactly as a quantization block splits `K`. A page table names
//! which physical page holds a sequence's logical one, and the walk routes `PAGE` to it.
//!
//! So the pool needs no ordering of its own, and the page size is the split's radix rather than a
//! number carried beside a lookup.
//!
//! The second half reads the same kind of pool with no walk over pages at all: the operand is
//! placed by its table ([`TileArgLaunch::indexed`] keyed by the sequence and its positions), the
//! kernel walks logical positions, and the table is resolved where each line's offset is taken.
#![allow(non_snake_case)]

use super::{Form, implied, uncut};
use crate::support::{HostData, HostDataType, TestInput};
use cubecl::std::tensor::TensorHandle;
use cubecl::{prelude::*, zspace::Shape};
use cubecl_tile::launch::{Bound, Delivery, Refusal};
use cubecl_tile::layout::{PageEntry, PhysicalAxisMap};
use cubecl_tile::*;

/// Which sequence.
const B: Axis = Axis(0);
/// Which of the sequence's own pages: what the table is addressed by.
const LP: Axis = Axis(1);
/// Which physical page of the pool: what the table names, and what the walk routes.
const PAGE: Axis = Axis(2);
/// Where inside a page.
const OFFSET: Axis = Axis(3);
/// The feature a slot carries.
const D: Axis = Axis(4);

const SEQS: usize = 2;
const PAGES_PER_SEQ: usize = 2;
const POOL_PAGES: usize = 4;
const PAGE_SIZE: usize = 2;
const FEATURES: usize = 2;
const SLOTS: usize = POOL_PAGES * PAGE_SIZE;

/// Page 2 is shared: sequence 0 holds it first, sequence 1 second. Nothing about the pool's order
/// says so, which is the point of a table.
const TABLE: [u32; SEQS * PAGES_PER_SEQ] = [2, 0, 3, 2];

/// `out[b] = Σ kv[slot]` over the slots of the pages sequence `b` holds.
#[cube(launch)]
fn paged_sum_kernel<E: Numeric>(
    kv: &TileArg<'_>,
    out: &TileArg<'_>,
    table: &Tensor<u32>,
    space: Partitioning,
    #[comptime] seq: Level,
    #[comptime] logical: Level,
    #[comptime] physical: Level,
    #[define(E)] _dtype: ElemType,
) {
    let kv = kv.tile(&space);
    let out = out.tile(&space);

    for sequence in space.over(&seq) {
        let b = sequence.coord(B);

        let mut total = out.at(&sequence);
        total.init::<E>(Monoid::identity::<E>(comptime!(Monoid::Sum)));

        for lp in sequence.over(&logical) {
            let phys = table[b * PAGES_PER_SEQ + lp.coord(LP)] as usize;

            // The same call MoE makes: the axis keeps the pool's four pages, the walk visits the
            // one the table named.
            for page in lp.over(&physical).routed(PAGE, phys) {
                let mut cell = out.at(&page);
                cell.reduce_axis_accumulate(&kv.at(&page), comptime!(Monoid::Sum));
            }
        }
    }
}

/// Every slot distinct, so reading a neighbouring page cannot land on the right answer.
fn kv_values() -> Vec<f32> {
    (0..SLOTS * FEATURES)
        .map(|i| {
            let (slot, d) = (i / FEATURES, i % FEATURES);
            (slot * 10 + d) as f32
        })
        .collect()
}

fn run(table: &[u32]) -> HostData {
    let client = cubecl::test_device().client();
    let f32_ty = f32::elem_type_native();
    let u32_ty = u32::elem_type_native();

    let launcher = implied(
        &client,
        Partitioning::new(
            Space::new(&[
                (B, SEQS),
                (LP, PAGES_PER_SEQ),
                (PAGE, POOL_PAGES),
                (OFFSET, PAGE_SIZE),
                (D, FEATURES),
            ]),
            Levels::leaf(&[(B, 1), (LP, 1), (PAGE, 1)])
                .walk_every(&[PAGE])
                .walk_every(&[LP])
                .walk_every(&[B])
                .build(),
        ),
        Form::Static,
    );

    let (kv_handle, _) = TestInput::builder(client.clone(), Shape::new([SLOTS, FEATURES]))
        .dtype(f32_ty)
        .custom(kv_values())
        .generate_with_f32_host_data();
    let table_handle = TestInput::builder(client.clone(), Shape::new([TABLE.len()]))
        .dtype(u32_ty)
        .custom(table.iter().map(|&p| p as f32).collect())
        .generate_without_host_data();
    let out_handle = TestInput::builder(client.clone(), Shape::new([SEQS, FEATURES]))
        .dtype(f32_ty)
        .custom(vec![-1.0; SEQS * FEATURES])
        .generate_without_host_data();

    paged_sum_kernel::launch(
        &client,
        launcher.cube_count(),
        launcher.cube_dim(),
        TileArgLaunch::new(
            kv_handle.binding().into_tensor_arg(),
            f32_ty,
            1,
            // The pool's one dimension read as two axes: `slot = page * PAGE_SIZE + offset`.
            TileSpec::new(Projection::new(
                &[PAGE, OFFSET, D],
                &[
                    PhysicalAxisMap::disjoint(&[(PAGE, PAGE_SIZE), (OFFSET, 1)]),
                    PhysicalAxisMap::of(D),
                ],
            )),
        ),
        TileArgLaunch::new(
            out_handle.clone().binding().into_tensor_arg(),
            f32_ty,
            1,
            TileSpec::direct(&[B, D]),
        ),
        table_handle.binding().into_tensor_arg(),
        launcher.partitioning_arg(),
        launcher.partitioning().level(0),
        launcher.partitioning().level(1),
        launcher.partitioning().level(2),
        f32_ty,
    );

    HostData::from_tensor_handle(&client, out_handle, HostDataType::F32)
}

/// The reference: each sequence over the slots of the pages its table names.
fn expected(table: &[u32]) -> Vec<Vec<f32>> {
    let values = kv_values();
    (0..SEQS)
        .map(|b| {
            (0..FEATURES)
                .map(|d| {
                    (0..PAGES_PER_SEQ)
                        .flat_map(|lp| {
                            let page = table[b * PAGES_PER_SEQ + lp] as usize;
                            (0..PAGE_SIZE).map(move |o| (page * PAGE_SIZE + o) * FEATURES + d)
                        })
                        .map(|i| values[i])
                        .sum()
                })
                .collect()
        })
        .collect()
}

/// Two page tables over the same pool: the answer follows the table, which it could not if the
/// route did nothing, since the walk would reduce the whole pool under both.
#[test]
fn a_routed_page_axis_reduces_each_sequence_over_the_pages_it_holds() {
    for table in [TABLE, [0, 1, 1, 0]] {
        let got = run(&table);
        let want = expected(&table);
        for (b, row) in want.iter().enumerate() {
            for (d, cell) in row.iter().enumerate() {
                assert_eq!(
                    got.get_f32(&[b, d]),
                    *cell,
                    "sequence {b} feature {d}: pages {:?}",
                    &table[b * PAGES_PER_SEQ..(b + 1) * PAGES_PER_SEQ]
                );
            }
        }
    }
}

/// The paged reads below: a position, as the sequences see it, and nothing of the pool.
const S: Axis = Axis(1);
/// The feature a position carries.
const F: Axis = Axis(2);

const SEQUENCES: usize = 3;
const PAGES_EACH: usize = 3;
const PAGE_LEN: usize = 2;
const POOL: usize = 7;
const FEATURES_READ: usize = 3;
/// Positions a step reads: more than a page, so every step but the last crosses one.
const READ_BLOCK: usize = 4;

/// Each sequence's pages, scattered over the pool; page 4 is sequence 2's twice, and page 0 is
/// held by two sequences.
const PAGES: [[u32; PAGES_EACH]; SEQUENCES] = [[5, 0, 3], [1, 6, 2], [4, 4, 0]];

/// `handle`, a `u32` page table `[sequences, pages_per_sequence]`, keyed by `keys`: the sequence
/// whose pages it lists, then the positions they hold.
fn page_table(handle: TensorHandle, keys: &[Axis]) -> TileArgLaunch {
    TileArgLaunch::new(
        handle.binding().into_tensor_arg(),
        u32::elem_type_native(),
        1,
        TileSpec::direct(keys),
    )
}

/// Which entry a read hands its cuts ([`Tile::with_page_entry`]).
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
enum Handing {
    /// None: the cuts read the table.
    Nothing,
    /// The entry of the step about to be cut.
    Own,
    /// The entry of the step after.
    Next,
    /// The entry of a step past the table.
    PastTable,
    /// The entry of the same step of the next sequence: the same column, another sequence.
    OtherSequence,
    /// The entry of the step being cut, read one step earlier and carried through the loop.
    Carried,
    /// An entry no table lists.
    Unlisted,
}

impl Handing {
    /// How far past the step being cut `Own`, `Next` and `PastTable` read their entry.
    fn steps_ahead(self) -> usize {
        match self {
            Handing::Nothing
            | Handing::Own
            | Handing::Unlisted
            | Handing::OtherSequence
            | Handing::Carried => 0,
            Handing::Next => 1,
            Handing::PastTable => PAGES_EACH * PAGE_LEN,
        }
    }
}

/// `out[b] = Σ pool[b's page of s, s % PAGE_LEN]` over positions `from..lens[b]` of sequence `b`,
/// read through the table its binding carries: the kernel walks logical positions, `read_block` of them
/// a step, and [`Tile::within`] cuts each sequence at its own length, handing each cut the entry
/// `handing` names.
#[cube(launch)]
fn paged_read_sum_kernel<E: Numeric>(
    pool: &TileArg<'_>,
    lens: &Tensor<u32>,
    out: &TileArg<'_>,
    #[comptime] from: usize,
    space: Partitioning,
    #[comptime] seq: Level,
    #[comptime] block: Level,
    #[comptime] read_block: usize,
    #[comptime] handing: Handing,
    #[define(E)] _dtype: ElemType,
) {
    let pool = pool.tile(&space);
    let out = out.tile(&space);

    for sequence in space.over(&seq) {
        let len = lens[sequence.coord(B)] as usize;

        let mut total = out.at(&sequence);
        total.init::<E>(Monoid::identity::<E>(comptime!(Monoid::Sum)));

        let positions = pool.within(S, from, len);
        let steps = sequence.over(&block);
        let mut carried = PageEntry::unlisted();
        if comptime!(handing == Handing::Carried) {
            carried = positions.read_page_entry(&steps.region(0));
        }
        for i in 0..len.saturating_sub(from).div_ceil(read_block) {
            let blk = steps.region(i);
            let mut cell = out.at(&blk);
            if comptime!(handing == Handing::Nothing) {
                cell.reduce_axis_accumulate(&positions.at(&blk), comptime!(Monoid::Sum));
            } else {
                let mut entry = PageEntry::unlisted();
                if comptime!(handing == Handing::Carried) {
                    entry = carried.clone();
                    carried = positions.read_page_entry(&steps.region(i + 1));
                } else if comptime!(handing == Handing::OtherSequence) {
                    let next = space.over(&seq).region((sequence.coord(B) + 1) % SEQUENCES);
                    entry = positions.read_page_entry(&next.over(&block).region(i));
                } else if comptime!(handing != Handing::Unlisted) {
                    let source = steps.region(i + comptime!(handing.steps_ahead()));
                    entry = positions.read_page_entry(&source);
                }
                let handed = positions.with_page_entry(&entry);
                cell.reduce_axis_accumulate(&handed.at(&blk), comptime!(Monoid::Sum));
            }
        }
    }
}

/// Every pool cell distinct, pages of `page_len` positions: `page·100 + row·10 + feature`.
fn pool_values(page_len: usize) -> Vec<f32> {
    (0..POOL * page_len * FEATURES_READ)
        .map(|i| {
            let (page, row, f) = (
                i / (page_len * FEATURES_READ),
                (i / FEATURES_READ) % page_len,
                i % FEATURES_READ,
            );
            (page * 100 + row * 10 + f) as f32
        })
        .collect()
}

fn run_paged_read(lens: &[u32; SEQUENCES]) -> HostData {
    run_paged_read_over(lens, &rows(&PAGES))
}

/// A read through [`PAGES`] `block` positions a step from position `from` on.
fn run_paged_read_from(lens: &[u32; SEQUENCES], block: usize, from: u32) -> HostData {
    run_paged_read_blocks(lens, &rows(&PAGES), block, from)
}

/// `pages` as a table of rows, any number of them and of any width.
fn rows(pages: &[[u32; PAGES_EACH]]) -> Vec<Vec<u32>> {
    pages.iter().map(|row| row.to_vec()).collect()
}

/// [`run_paged_read`] through `pages`, whose rows and columns need not be the launch's: the table
/// sets the operand's extent, and the launch walks its own.
fn run_paged_read_over(lens: &[u32; SEQUENCES], pages: &[Vec<u32>]) -> HostData {
    run_paged_read_blocks(lens, pages, READ_BLOCK, 0)
}

/// [`run_paged_read_over`] `block` positions a step from position `from` on.
fn run_paged_read_blocks(
    lens: &[u32; SEQUENCES],
    pages: &[Vec<u32>],
    block: usize,
    from: u32,
) -> HostData {
    let columns = pages.first().map_or(PAGES_EACH, Vec::len);
    run_paged_read_table(
        lens,
        &[pages.len(), columns],
        pages.iter().flatten().copied().collect(),
        PAGE_LEN,
        block,
        from,
        Handing::Nothing,
        TileArgLaunch::indexed,
    )
}

/// How a read places its pool by its table: [`TileArgLaunch::indexed`] or
/// [`TileArgLaunch::indexed_covering`].
type Bind = fn(TileArgLaunch, TileArgLaunch, Indexing) -> Result<TileArgLaunch, Refusal>;

/// [`run_paged_read_blocks`] through a table of `entries` laid out as `shape`, any rank, the pool
/// bound by `bind`.
#[allow(clippy::too_many_arguments)]
fn run_paged_read_table(
    lens: &[u32; SEQUENCES],
    shape: &[usize],
    entries: Vec<u32>,
    page_len: usize,
    block: usize,
    from: u32,
    handing: Handing,
    bind: Bind,
) -> HostData {
    let client = cubecl::test_device().client();
    let f32_ty = f32::elem_type_native();
    let u32_ty = u32::elem_type_native();

    let launcher = implied(
        &client,
        Partitioning::new(
            Space::new(&[
                (B, SEQUENCES),
                (
                    S,
                    shape.get(1).copied().unwrap_or(0).max(PAGES_EACH) * page_len,
                ),
                (F, FEATURES_READ),
            ]),
            Levels::leaf(&[(B, 1), (S, block)])
                .walk_every(&[S])
                .walk_every(&[B])
                .build(),
        ),
        Form::Static,
    );

    let (pool_handle, _) =
        TestInput::builder(client.clone(), Shape::new([POOL, page_len, FEATURES_READ]))
            .dtype(f32_ty)
            .custom(pool_values(page_len))
            .generate_with_f32_host_data();
    let table_handle = TestInput::builder(client.clone(), Shape::from(shape.to_vec()))
        .dtype(u32_ty)
        .custom(entries.iter().map(|&p| p as f32).collect())
        .generate_without_host_data();
    let lens_handle = TestInput::builder(client.clone(), Shape::new([SEQUENCES]))
        .dtype(u32_ty)
        .custom(lens.iter().map(|&l| l as f32).collect())
        .generate_without_host_data();
    let out_handle = TestInput::builder(client.clone(), Shape::new([SEQUENCES, FEATURES_READ]))
        .dtype(f32_ty)
        .custom(vec![-1.0; SEQUENCES * FEATURES_READ])
        .generate_without_host_data();

    let pool = launcher
        .arg(pool_handle.binding(), f32_ty)
        .axes(&[B, S, F])
        .build()
        .unwrap()
        .arg();
    let pool = bind(
        pool,
        page_table(table_handle, &[B, S]),
        Indexing::Coordinate(B),
    )
    .unwrap();

    paged_read_sum_kernel::launch(
        &client,
        launcher.cube_count(),
        launcher.cube_dim(),
        pool,
        lens_handle.binding().into_tensor_arg(),
        TileArgLaunch::new(
            out_handle.clone().binding().into_tensor_arg(),
            f32_ty,
            1,
            TileSpec::direct(&[B, F]),
        ),
        from as usize,
        launcher.partitioning_arg(),
        launcher.partitioning().level(0),
        launcher.partitioning().level(1),
        block,
        handing,
        f32_ty,
    );

    HostData::from_tensor_handle(&client, out_handle, HostDataType::F32)
}

/// [`paged_read_sum_kernel`] walking each sequence a page at a time and the page's positions
/// `step` at a time inside it: each page's window is held and its steps cut from it. A `from` of
/// 0 keeps the alignment the levels prove; any other is a runtime shift that loses it.
#[cube(launch)]
fn paged_page_walk_sum_kernel<E: Numeric>(
    pool: &TileArg<'_>,
    lens: &Tensor<u32>,
    out: &TileArg<'_>,
    #[comptime] from: usize,
    space: Partitioning,
    #[comptime] seq: Level,
    #[comptime] page: Level,
    #[comptime] page_len: usize,
    #[define(E)] _dtype: ElemType,
) {
    let pool = pool.tile(&space);
    let out = out.tile(&space);

    for sequence in space.over(&seq) {
        let len = lens[sequence.coord(B)] as usize;

        let mut total = out.at(&sequence);
        total.init::<E>(Monoid::identity::<E>(comptime!(Monoid::Sum)));

        let positions = pool.within(S, from, len);
        for held in sequence
            .over(&page)
            .range(0, len.saturating_sub(from).div_ceil(page_len))
        {
            let page_window = positions.at(&held);
            for step in held.walk() {
                let mut cell = out.at(&step);
                cell.reduce_axis_accumulate(&page_window.at(&step), comptime!(Monoid::Sum));
            }
        }
    }
}

/// Positions a page holds in the page walk: several steps of more than one position fit in it.
const WALK_PAGE_LEN: usize = 4;

/// [`paged_page_walk_sum_kernel`] over [`PAGES`] of [`WALK_PAGE_LEN`] positions, `step` positions
/// a step inside each page, from position `from` on.
fn run_paged_page_walk(lens: &[u32; SEQUENCES], step: usize, from: u32) -> HostData {
    let client = cubecl::test_device().client();
    let f32_ty = f32::elem_type_native();
    let u32_ty = u32::elem_type_native();

    let launcher = implied(
        &client,
        Partitioning::new(
            Space::new(&[
                (B, SEQUENCES),
                (S, PAGES_EACH * WALK_PAGE_LEN),
                (F, FEATURES_READ),
            ]),
            Levels::leaf(&[(B, 1), (S, step)])
                .walk(&[(S, WALK_PAGE_LEN / step)])
                .walk_every(&[S])
                .walk_every(&[B])
                .build(),
        ),
        Form::Static,
    );

    let (pool_handle, _) = TestInput::builder(
        client.clone(),
        Shape::new([POOL, WALK_PAGE_LEN, FEATURES_READ]),
    )
    .dtype(f32_ty)
    .custom(pool_values(WALK_PAGE_LEN))
    .generate_with_f32_host_data();
    let table_handle = TestInput::builder(client.clone(), Shape::new([SEQUENCES, PAGES_EACH]))
        .dtype(u32_ty)
        .custom(PAGES.iter().flatten().map(|&p| p as f32).collect())
        .generate_without_host_data();
    let lens_handle = TestInput::builder(client.clone(), Shape::new([SEQUENCES]))
        .dtype(u32_ty)
        .custom(lens.iter().map(|&l| l as f32).collect())
        .generate_without_host_data();
    let out_handle = TestInput::builder(client.clone(), Shape::new([SEQUENCES, FEATURES_READ]))
        .dtype(f32_ty)
        .custom(vec![-1.0; SEQUENCES * FEATURES_READ])
        .generate_without_host_data();

    let pool = launcher
        .arg(pool_handle.binding(), f32_ty)
        .axes(&[B, S, F])
        .build()
        .unwrap()
        .arg()
        .indexed_covering(page_table(table_handle, &[B, S]), Indexing::Coordinate(B))
        .unwrap();

    paged_page_walk_sum_kernel::launch(
        &client,
        launcher.cube_count(),
        launcher.cube_dim(),
        pool,
        lens_handle.binding().into_tensor_arg(),
        TileArgLaunch::new(
            out_handle.clone().binding().into_tensor_arg(),
            f32_ty,
            1,
            TileSpec::direct(&[B, F]),
        ),
        from as usize,
        launcher.partitioning_arg(),
        launcher.partitioning().level(0),
        launcher.partitioning().level(1),
        WALK_PAGE_LEN,
        f32_ty,
    );

    HostData::from_tensor_handle(&client, out_handle, HostDataType::F32)
}

/// The reference: each sequence's positions looked up in its row of the table on the host.
fn expected_paged_read(lens: &[u32; SEQUENCES]) -> Vec<Vec<f32>> {
    expected_paged_read_over(lens, &rows(&PAGES))
}

/// [`expected_paged_read`] through `pages`, an entry past the pool reading its last page and a
/// position past the table reading zero.
fn expected_paged_read_over(lens: &[u32; SEQUENCES], pages: &[Vec<u32>]) -> Vec<Vec<f32>> {
    expected_paged_read_from(lens, pages, 0, PAGE_LEN)
}

/// [`expected_paged_read_over`] from position `from` on, over pages of `page_len` positions.
fn expected_paged_read_from(
    lens: &[u32; SEQUENCES],
    pages: &[Vec<u32>],
    from: u32,
    page_len: usize,
) -> Vec<Vec<f32>> {
    let values = pool_values(page_len);
    (0..SEQUENCES)
        .map(|b| {
            (0..FEATURES_READ)
                .map(|f| {
                    (from as usize..lens[b] as usize)
                        .map(|s| {
                            let page = pages.get(b).and_then(|row| row.get(s / page_len));
                            page.map_or(0.0, |&page| {
                                let page = (page as usize).min(POOL - 1);
                                values[(page * page_len + s % page_len) * FEATURES_READ + f]
                            })
                        })
                        .sum()
                })
                .collect()
        })
        .collect()
}

/// A sequence read through its page table reduces the positions its pages hold, in its table's
/// order: steps that cross pages, a length that ends inside a page, a full sequence and one of a
/// single position. Reading the pool where it lies would reduce pages 0, 1 and 2 for sequence 0.
#[test]
fn a_paged_operand_reads_each_sequence_through_its_table() {
    for lens in [[5, 6, 1], [6, 3, 4]] {
        let got = run_paged_read(&lens);
        let want = expected_paged_read(&lens);
        for (b, row) in want.iter().enumerate() {
            for (f, cell) in row.iter().enumerate() {
                assert_eq!(
                    got.get_f32(&[b, f]),
                    *cell,
                    "sequence {b} feature {f}: {} positions of pages {:?}",
                    lens[b],
                    PAGES[b]
                );
            }
        }
    }
}

/// A table entry past the pool reads its last page rather than past the buffer: a table is device
/// data a launch may never have written.
#[test]
fn a_paged_read_stays_in_the_pool_and_the_table() {
    let pages = rows(&[[5, 99, 3], [1, 6, 2], [4, u32::MAX, 0]]);
    let lens = [6, 3, 5];
    let got = run_paged_read_over(&lens, &pages);
    let want = expected_paged_read_over(&lens, &pages);
    for (b, row) in want.iter().enumerate() {
        for (f, cell) in row.iter().enumerate() {
            assert_eq!(got.get_f32(&[b, f]), *cell, "sequence {b} feature {f}");
        }
    }
}

/// Steps of a page or less resolve both pages a window can touch as it is cut: aligned to the
/// pages, inside one, and straddling two from an odd start.
#[test]
fn a_paged_read_of_steps_within_a_page_reads_each_sequence_through_its_table() {
    let lens = [6, 5, 3];
    for (block, from) in [(1, 0), (2, 0), (2, 1), (1, 3)] {
        let got = run_paged_read_from(&lens, block, from);
        let want = expected_paged_read_from(&lens, &rows(&PAGES), from, PAGE_LEN);
        for (b, row) in want.iter().enumerate() {
            for (f, cell) in row.iter().enumerate() {
                assert_eq!(
                    got.get_f32(&[b, f]),
                    *cell,
                    "{block} a step from {from}: sequence {b} feature {f}"
                );
            }
        }
    }
}

/// Steps wider than a page pin every page they may reach: three pages from an unaligned or an
/// aligned origin, four with the last pinned column past the table, and past the pin limit all
/// the lines resolved one by one.
#[test]
fn a_paged_read_of_steps_wider_than_a_page_reads_each_sequence_through_its_table() {
    let lens = [6, 5, 3];
    for (block, from) in [(READ_BLOCK, 1), (6, 0), (6, 1), (4, 3)] {
        let got = run_paged_read_from(&lens, block, from);
        let want = expected_paged_read_from(&lens, &rows(&PAGES), from, PAGE_LEN);
        for (b, row) in want.iter().enumerate() {
            for (f, cell) in row.iter().enumerate() {
                assert_eq!(
                    got.get_f32(&[b, f]),
                    *cell,
                    "{block} a step from {from}: sequence {b} feature {f}"
                );
            }
        }
    }
    let pages = vec![
        vec![5, 0, 3, 6, 1, 2],
        vec![1, 6, 2, 4, 0, 5],
        vec![4, 4, 0, 3, 2, 6],
    ];
    let lens = [12, 9, 7];
    for (block, from) in [(8, 0), (8, 1)] {
        let got = run_paged_read_blocks(&lens, &pages, block, from);
        let want = expected_paged_read_from(&lens, &pages, from, PAGE_LEN);
        for (b, row) in want.iter().enumerate() {
            for (f, cell) in row.iter().enumerate() {
                assert_eq!(
                    got.get_f32(&[b, f]),
                    *cell,
                    "{block} a step from {from}: sequence {b} feature {f}"
                );
            }
        }
    }
}

/// Steps cut straight from the whole sequence, over pages of four positions: a step of 1, 2 or 4
/// has an origin aligned to its own size, which proves it inside one page; a step of 3 proves
/// nothing and may straddle two.
#[test]
fn a_flat_walk_over_larger_pages_reads_each_sequence_through_its_table() {
    let entries = PAGES.iter().flatten().copied().collect::<Vec<_>>();
    let lens = [9, 12, 7];
    for (block, from) in [(1, 0), (2, 0), (3, 0), (4, 0), (2, 1), (4, 3)] {
        let got = run_paged_read_table(
            &lens,
            &[SEQUENCES, PAGES_EACH],
            entries.clone(),
            WALK_PAGE_LEN,
            block,
            from,
            Handing::Nothing,
            TileArgLaunch::indexed,
        );
        let want = expected_paged_read_from(&lens, &rows(&PAGES), from, WALK_PAGE_LEN);
        for (b, row) in want.iter().enumerate() {
            for (f, cell) in row.iter().enumerate() {
                assert_eq!(
                    got.get_f32(&[b, f]),
                    *cell,
                    "steps of {block} from {from}: sequence {b} feature {f}"
                );
            }
        }
    }
}

/// A cut handed a page entry reads what it reads without one, whichever entry that is, including
/// one of another sequence and one carried from the step before. Steps of one position, of a page
/// and of two pages, so the cut is inside a page and across two, and a page-sized step from an odd
/// start, whose window straddles two pages.
#[test]
fn a_cut_handed_a_page_entry_reads_each_sequence_through_its_table() {
    let entries = PAGES.iter().flatten().copied().collect::<Vec<_>>();
    let lens = [6, 5, 3];
    for handing in [
        Handing::Own,
        Handing::Next,
        Handing::PastTable,
        Handing::Unlisted,
        Handing::OtherSequence,
        Handing::Carried,
    ] {
        for (block, from) in [(1, 0), (PAGE_LEN, 0), (READ_BLOCK, 0), (PAGE_LEN, 1)] {
            let got = run_paged_read_table(
                &lens,
                &[SEQUENCES, PAGES_EACH],
                entries.clone(),
                PAGE_LEN,
                block,
                from,
                handing,
                TileArgLaunch::indexed,
            );
            let want = expected_paged_read_from(&lens, &rows(&PAGES), from, PAGE_LEN);
            for (b, row) in want.iter().enumerate() {
                for (f, cell) in row.iter().enumerate() {
                    assert_eq!(
                        got.get_f32(&[b, f]),
                        *cell,
                        "{handing:?}, steps of {block} from {from}: sequence {b} feature {f}"
                    );
                }
            }
        }
    }
}

/// A sequence walked a page at a time reads each page through the table once and every step
/// inside it through that page: steps of one position, of two and of the whole page, lengths
/// ending inside a page, over pages scattered and shared across sequences. Started past 0, each
/// page's window straddles two pages, so its steps cannot keep its pin.
#[test]
fn a_sequence_walked_a_page_at_a_time_reads_each_sequence_through_its_table() {
    for (step, from) in [
        (1, 0),
        (2, 0),
        (WALK_PAGE_LEN, 0),
        (2, 1),
        (1, 3),
        (WALK_PAGE_LEN, 3),
    ] {
        for lens in [[9, 12, 1], [12, 6, 7]] {
            let got = run_paged_page_walk(&lens, step, from);
            let want = expected_paged_read_from(&lens, &rows(&PAGES), from, WALK_PAGE_LEN);
            for (b, row) in want.iter().enumerate() {
                for (f, cell) in row.iter().enumerate() {
                    assert_eq!(
                        got.get_f32(&[b, f]),
                        *cell,
                        "steps of {step} from {from}: sequence {b} feature {f}, {} positions",
                        lens[b]
                    );
                }
            }
        }
    }
}

/// The table sets the operand's extent, not the launch: a table of fewer rows or narrower rows
/// than the launch walks reads zero past them, down to a table of one entry. That no load lands
/// past the table is `PageTable::lookup`'s clamp, which values cannot show: the mask drops
/// whatever such a load returns.
#[test]
fn a_paged_read_past_the_table_reads_zero() {
    let lens = [6, 5, 4];
    for pages in [
        rows(&PAGES[..2]),
        PAGES.iter().map(|row| row[..2].to_vec()).collect(),
        vec![vec![PAGES[0][0]]],
    ] {
        let got = run_paged_read_over(&lens, &pages);
        let want = expected_paged_read_over(&lens, &pages);
        for (b, row) in want.iter().enumerate() {
            for (f, cell) in row.iter().enumerate() {
                assert_eq!(
                    got.get_f32(&[b, f]),
                    *cell,
                    "sequence {b} feature {f} through {pages:?}"
                );
            }
        }
    }
}

/// A table stated to cover the launch reads as one checked against its shape wherever the
/// launch keeps its reads in bounds: steps that cross pages, lengths ending inside a page.
#[test]
fn a_covering_table_reads_each_sequence_through_it() {
    let entries = PAGES.iter().flatten().copied().collect::<Vec<_>>();
    for lens in [[5, 6, 1], [6, 3, 4]] {
        let shape = [SEQUENCES, PAGES_EACH];
        let got = run_paged_read_table(
            &lens,
            &shape,
            entries.clone(),
            PAGE_LEN,
            READ_BLOCK,
            0,
            Handing::Nothing,
            TileArgLaunch::indexed_covering,
        );
        for (b, row) in expected_paged_read(&lens).iter().enumerate() {
            for (f, cell) in row.iter().enumerate() {
                assert_eq!(got.get_f32(&[b, f]), *cell, "sequence {b} feature {f}");
            }
        }
    }
}

/// What a kernel could not read through a table is refused as the table is bound: a table of
/// another rank than its two keys, one dim for both axes, pages along the innermost dim, a pool of
/// no pages, and a delivery that moves
/// runs.
#[test]
fn a_paged_operand_the_table_cannot_resolve_is_refused_at_build() {
    let client = cubecl::test_device().client();
    let launcher = implied(
        &client,
        Partitioning::new(
            Space::new(&[(B, SEQUENCES), (S, PAGES_EACH * 3), (F, 4)]),
            Levels::leaf(&[(B, 1), (S, READ_BLOCK)])
                .walk_every(&[S])
                .walk_every(&[B])
                .build(),
        ),
        Form::Static,
    );
    let pool = |shape: [usize; 3]| {
        TestInput::builder(client.clone(), Shape::new(shape))
            .dtype(f32::elem_type_native())
            .zeros()
            .generate_without_host_data()
    };
    let table = TestInput::builder(client.clone(), Shape::new([SEQUENCES, PAGES_EACH]))
        .dtype(u32::elem_type_native())
        .zeros()
        .generate_without_host_data();
    let placed = |bound: Result<Bound, Refusal>, (sequence, position)| {
        bound
            .and_then(|pool| {
                pool.arg().indexed(
                    page_table(table.clone(), &[sequence, position]),
                    Indexing::Coordinate(sequence),
                )
            })
            .err()
    };
    let refusal = |binding, axes: &[Axis], keys, width| {
        let pool = launcher
            .arg(binding, f32::elem_type_native())
            .axes(axes)
            .vectorize(width)
            .build();
        placed(pool, keys)
    };
    let flat = TestInput::builder(client.clone(), Shape::new([SEQUENCES * PAGES_EACH]))
        .dtype(u32::elem_type_native())
        .zeros()
        .generate_without_host_data();
    let rows = pool([POOL, 3, 4]);
    let flat_table = launcher
        .arg(rows.binding(), f32::elem_type_native())
        .axes(&[B, S, F])
        .build()
        .and_then(|pool| {
            pool.arg()
                .indexed(page_table(flat, &[B, S]), Indexing::Coordinate(B))
        })
        .err();
    assert_eq!(
        flat_table,
        Some(Refusal::IndexTableRank { keys: 2, rank: 1 })
    );
    let a = pool([POOL, 3, 4]);
    assert_eq!(
        refusal(a.binding(), &[B, S, F], (S, S), 1),
        Some(Refusal::PagingOfOneDim(S))
    );
    let b = pool([POOL, 3, 4]);
    assert_eq!(
        refusal(b.binding(), &[S, F, B], (B, S), 1),
        Some(Refusal::PagesAlongALine(B))
    );
    let e = pool([0, 3, 4]);
    assert_eq!(
        refusal(e.binding(), &[B, S, F], (B, S), 1),
        Some(Refusal::PagingOfAnEmptyPool)
    );
    let f = pool([POOL, 3, 4]);
    let by_runs = placed(
        launcher
            .arg(f.binding(), f32::elem_type_native())
            .axes(&[B, S, F])
            .delivery(Delivery::AsyncBulk)
            .build(),
        (B, S),
    );
    assert_eq!(by_runs, Some(Refusal::PagingMovedByRuns));
}

/// The pool is bound as it lies, so where its pages sit among its dims is no matter: pages inner
/// to a head dim, and a head dim of one whose stride is anything, as an `unsqueeze` leaves it.
#[test]
fn a_paged_pool_binds_wherever_its_pages_lie() {
    const H: Axis = Axis(3);
    let client = cubecl::test_device().client();
    let launcher = implied(
        &client,
        Partitioning::new(
            Space::new(&[(B, SEQUENCES), (S, PAGES_EACH * PAGE_LEN), (F, 4), (H, 2)]),
            Levels::leaf(&[(B, 1), (S, READ_BLOCK)])
                .walk_every(&[S])
                .walk_every(&[B])
                .build(),
        ),
        Form::Static,
    );
    let pool = TestInput::builder(client.clone(), Shape::new([2, POOL, PAGE_LEN, 4]))
        .dtype(f32::elem_type_native())
        .zeros()
        .generate_without_host_data();
    let table = TestInput::builder(client.clone(), Shape::new([SEQUENCES, PAGES_EACH]))
        .dtype(u32::elem_type_native())
        .zeros()
        .generate_without_host_data();
    let placed = |bound: Bound| {
        bound
            .arg()
            .indexed(page_table(table.clone(), &[B, S]), Indexing::Coordinate(B))
    };
    let heads_outside = launcher
        .arg(pool.binding(), f32::elem_type_native())
        .axes(&[H, B, S, F])
        .build()
        .and_then(placed);
    assert!(heads_outside.is_ok(), "{:?}", heads_outside.err());
    let unsqueezed = TestInput::builder(client.clone(), Shape::new([POOL, PAGE_LEN, 4]))
        .dtype(f32::elem_type_native())
        .zeros()
        .generate_without_host_data();
    let mut binding = unsqueezed.binding();
    binding.shape = [POOL, 1, PAGE_LEN, 4].into();
    binding.strides = [PAGE_LEN * 4, 4096, 4, 1].into();
    let one_head = launcher
        .arg(binding, f32::elem_type_native())
        .axes(&[B, H, S, F])
        .build()
        .and_then(placed);
    assert!(one_head.is_ok(), "{:?}", one_head.err());
}

/// `out = pool` over the sequences, every scale it carries applied, the pool and its scales read
/// through the tables their bindings carry.
#[cube(launch)]
fn paged_dequantize_kernel(pool: &TileArg<'_>, out: &TileArg<'_>, space: Partitioning) {
    let pool = pool.tile(&space);
    let mut out = out.tile(&space);
    out.dequantize_from(&pool);
}

/// The pools [`run_packed_paged`] reads: a pool of values, `q8` four to a word along the features
/// where `packed`, and where `scaled`, a pool of a scale per position read through the same table.
#[derive(Clone, Copy, Debug)]
struct PagedPools {
    packed: bool,
    scaled: bool,
}

/// `value(page, row, feature)`: distinct, inside a signed byte.
fn packed_value(page: usize, row: usize, f: usize) -> i32 {
    (page * 16 + row * 4 + f) as i32 - 64
}

/// The scale of a position of the pool.
fn position_scale(page: usize, row: usize) -> f32 {
    0.5 * (1 + (page + row) % 3) as f32
}

/// The pools `pools` names copied out as the sequences [`PAGES`] lays out; `None` where the device
/// can't serve a whole word a line.
fn run_packed_paged(pools: PagedPools) -> Option<HostData> {
    let field = cubecl::quant::scheme::QuantValue::Q8S;
    let bits = field.size_bits();
    let features = 32 / bits;
    let client = cubecl::test_device().client();
    if features > client.properties().hardware.max_vector_size {
        return None;
    }
    let (f32_ty, u32_ty) = (f32::elem_type_native(), u32::elem_type_native());
    let positions = PAGES_EACH * PAGE_LEN;

    let cells = (0..POOL * PAGE_LEN).map(|cell| (cell / PAGE_LEN, cell % PAGE_LEN));
    let space = Space::new(&[(B, SEQUENCES), (S, positions), (F, features)]);
    let launcher = uncut(&client, &space, &space);
    let table = TestInput::builder(client.clone(), Shape::new([SEQUENCES, PAGES_EACH]))
        .dtype(u32_ty)
        .custom(PAGES.iter().flatten().map(|&p| p as f32).collect())
        .generate_without_host_data();
    let out = TestInput::builder(client.clone(), Shape::new([SEQUENCES, positions, features]))
        .dtype(f32_ty)
        .zeros()
        .generate_without_host_data();

    let values = match pools.packed {
        true => {
            let words: Vec<u32> = cells
                .clone()
                .map(|(page, row)| {
                    (0..features).fold(0u32, |word, f| {
                        let v = packed_value(page, row, f) as u32 & 0xFF;
                        word | (v << (f * bits))
                    })
                })
                .collect();
            // Shape and strides count values; the packing says how many share a stored word.
            let pool = cubecl::std::tensor::TensorHandle::new_contiguous(
                vec![POOL, PAGE_LEN, features],
                client.create(cubecl::bytes::Bytes::from_elems(words)),
                u32_ty,
            );
            launcher
                .arg(pool.binding(), u32_ty)
                .axes(&[B, S, F])
                .vectorize(features)
                .packed(field)
        }
        false => {
            let (pool, _) =
                TestInput::builder(client.clone(), Shape::new([POOL, PAGE_LEN, features]))
                    .dtype(f32_ty)
                    .custom(
                        cells
                            .clone()
                            .flat_map(|(page, row)| {
                                (0..features).map(move |f| packed_value(page, row, f) as f32)
                            })
                            .collect(),
                    )
                    .generate_with_f32_host_data();
            launcher
                .arg(pool.binding(), f32_ty)
                .axes(&[B, S, F])
                .vectorize(features)
        }
    }
    .build()
    .unwrap()
    .arg();
    let values = match pools.scaled {
        true => {
            let (scale_pool, _) = TestInput::builder(client.clone(), Shape::new([POOL, PAGE_LEN]))
                .dtype(f32_ty)
                .custom(cells.map(|(page, row)| position_scale(page, row)).collect())
                .generate_with_f32_host_data();
            let scales = launcher
                .arg(scale_pool.binding(), f32_ty)
                .axes(&[B, S])
                .build()
                .unwrap()
                .arg()
                .indexed(page_table(table.clone(), &[B, S]), Indexing::Coordinate(B))
                .unwrap();
            values.scales(scales)
        }
        false => values,
    };
    let values = values
        .indexed(page_table(table, &[B, S]), Indexing::Coordinate(B))
        .unwrap();
    let out_arg = launcher
        .arg(out.clone().binding(), f32_ty)
        .axes(&[B, S, F])
        .vectorize(features)
        .build()
        .unwrap();

    paged_dequantize_kernel::launch(
        &client,
        launcher.cube_count(),
        launcher.cube_dim(),
        values,
        out_arg.arg(),
        launcher.partitioning_arg(),
    );

    Some(HostData::from_tensor_handle(
        &client,
        out,
        HostDataType::F32,
    ))
}

/// [`run_packed_paged`] checked against the pool looked up in [`PAGES`] on the host.
fn check_packed_paged(pools: PagedPools) {
    let Some(got) = run_packed_paged(pools) else {
        crate::support::ValidationResult::Skipped(
            "device vectors cap below a four-value word".to_string(),
        )
        .enforce();
        return;
    };
    for (b, row) in PAGES.iter().enumerate() {
        for s in 0..PAGES_EACH * PAGE_LEN {
            let (page, at) = (row[s / PAGE_LEN] as usize, s % PAGE_LEN);
            let scale = match pools.scaled {
                true => position_scale(page, at),
                false => 1.0,
            };
            for f in 0..4 {
                let want = packed_value(page, at, f) as f32 * scale;
                let have = got.get_f32(&[b, s, f]);
                assert_eq!(
                    have, want,
                    "{pools:?}: sequence {b} position {s} feature {f}"
                );
            }
        }
    }
}

/// A packed pool is paged like a plain one: its words lie along the features, the innermost dim,
/// across no page, so the table resolves them as it does a plain pool's lines.
#[test]
fn a_packed_pool_reads_each_sequence_through_its_table() {
    check_packed_paged(PagedPools {
        packed: true,
        scaled: false,
    });
}

/// The scales of a paged packed pool are a coarser operand bound as a pool of its own, read
/// through the same table: each rides the values, and each reads the pages its sequence names.
#[test]
fn a_packed_pool_under_paged_scales_reads_each_sequence_through_its_table() {
    check_packed_paged(PagedPools {
        packed: true,
        scaled: true,
    });
}