span-02: single-source release folder (engine final path, soft blend)

19886eb verified 8 days ago

2.88 kB

//! Launchers of cuda/spec.cu: the device side of speculative decoding (verify-batch metadata, greedy acceptance, draft input / pick, hidden-state pool). Pointers are u64 device addresses.
use crate::backend::Arg::{F, I, P};
use crate::ops::Gpu;
use anyhow::Result;

pub const KMAX: usize = 8;
/// ints per (cycle, sequence) emit record: [n_emit, accepted, done, tok_0 .. tok_KMAX]
pub const EMIT_W: usize = 3 + KMAX + 1;

impl Gpu {
    #[allow(clippy::too_many_arguments)]
    pub fn spec_build_verify(&self, tokens: u64, pos: u64, slot: u64, bt_off: u64, kv_hi: u64, kv_hi_sl: u64, seq_bt: u64, x0: u64, len: u64, drafts: u64, b: usize, k: usize, maxp: usize) -> Result<()> {
        assert!(k < KMAX);
        self.be.launch("spec", "spec_build_verify", [b as u32, 1, 1], [32, 1, 1], &[P(tokens), P(pos), P(slot), P(bt_off), P(kv_hi), P(kv_hi_sl), P(seq_bt), P(x0), P(len), P(drafts), I(k as i32), I(maxp as i32)])
    }
    #[allow(clippy::too_many_arguments)]
    pub fn spec_accept(&self, emit: u64, ctr: u64, ring: usize, len: u64, x0: u64, gen: u64, done: u64, h0: u64, g: u64, drafts: u64, hidden: u64, stop_ids: u64, maxgen: usize, b: usize, k: usize, hb: usize) -> Result<()> {
        self.be.launch(
            "spec",
            "spec_accept",
            [b as u32, 1, 1],
            [256, 1, 1],
            &[P(emit), P(ctr), I(ring as i32), P(len), P(x0), P(gen), P(done), P(h0), P(g), P(drafts), P(hidden), P(stop_ids), I(maxgen as i32), I(b as i32), I(k as i32), I(hb as i32)],
        )
    }
    /// x[s] = concat(embed(tok[s*tok_stride]) * scale, hs[s]) with the target's fp8 tied table
    #[allow(clippy::too_many_arguments)]
    pub fn spec_draft_input(&self, x: u64, table: u64, row_scale: u64, tok: u64, tok_stride: usize, hs: u64, b: usize, hb: usize, scale: f32) -> Result<()> {
        self.be.launch("spec", "spec_draft_input_fp8", [b as u32, 1, 1], [256, 1, 1], &[P(x), P(table), P(row_scale), P(tok), I(tok_stride as i32), P(hs), I(hb as i32), F(scale)])
    }
    pub fn spec_draft_pick(&self, drafts: u64, idx: u64, map: u64, j: usize, b: usize) -> Result<()> {
        self.be.launch("spec", "spec_draft_pick", [1, 1, 1], [b as u32, 1, 1], &[P(drafts), P(idx), P(map), I(j as i32)])
    }
    pub fn spec_draft_pos(&self, dpos: u64, len: u64, b: usize) -> Result<()> {
        self.be.launch("spec", "spec_draft_pos", [1, 1, 1], [b as u32, 1, 1], &[P(dpos), P(len)])
    }
    pub fn spec_h_gather(&self, h0: u64, hpool: u64, slot: u64, b: usize, hb: usize) -> Result<()> {
        self.be.launch("spec", "spec_h_gather", [b as u32, 1, 1], [256, 1, 1], &[P(h0), P(hpool), P(slot), I(hb as i32)])
    }
    pub fn spec_h_scatter(&self, hpool: u64, h0: u64, slot: u64, b: usize, hb: usize) -> Result<()> {
        self.be.launch("spec", "spec_h_scatter", [b as u32, 1, 1], [256, 1, 1], &[P(hpool), P(h0), P(slot), I(hb as i32)])
    }
}

	//! Launchers of cuda/spec.cu: the device side of speculative decoding (verify-batch metadata, greedy acceptance, draft input / pick, hidden-state pool). Pointers are u64 device addresses.
	use crate::backend::Arg::{F, I, P};
	use crate::ops::Gpu;
	use anyhow::Result;

	pub const KMAX: usize = 8;
	/// ints per (cycle, sequence) emit record: [n_emit, accepted, done, tok_0 .. tok_KMAX]
	pub const EMIT_W: usize = 3 + KMAX + 1;

	impl Gpu {
	#[allow(clippy::too_many_arguments)]
	pub fn spec_build_verify(&self, tokens: u64, pos: u64, slot: u64, bt_off: u64, kv_hi: u64, kv_hi_sl: u64, seq_bt: u64, x0: u64, len: u64, drafts: u64, b: usize, k: usize, maxp: usize) -> Result<()> {
	assert!(k < KMAX);
	self.be.launch("spec", "spec_build_verify", [b as u32, 1, 1], [32, 1, 1], &[P(tokens), P(pos), P(slot), P(bt_off), P(kv_hi), P(kv_hi_sl), P(seq_bt), P(x0), P(len), P(drafts), I(k as i32), I(maxp as i32)])
	}
	#[allow(clippy::too_many_arguments)]
	pub fn spec_accept(&self, emit: u64, ctr: u64, ring: usize, len: u64, x0: u64, gen: u64, done: u64, h0: u64, g: u64, drafts: u64, hidden: u64, stop_ids: u64, maxgen: usize, b: usize, k: usize, hb: usize) -> Result<()> {
	self.be.launch(
	"spec",
	"spec_accept",
	[b as u32, 1, 1],
	[256, 1, 1],
	&[P(emit), P(ctr), I(ring as i32), P(len), P(x0), P(gen), P(done), P(h0), P(g), P(drafts), P(hidden), P(stop_ids), I(maxgen as i32), I(b as i32), I(k as i32), I(hb as i32)],
	)
	}
	/// x[s] = concat(embed(tok[stok_stride]) scale, hs[s]) with the target's fp8 tied table
	#[allow(clippy::too_many_arguments)]
	pub fn spec_draft_input(&self, x: u64, table: u64, row_scale: u64, tok: u64, tok_stride: usize, hs: u64, b: usize, hb: usize, scale: f32) -> Result<()> {
	self.be.launch("spec", "spec_draft_input_fp8", [b as u32, 1, 1], [256, 1, 1], &[P(x), P(table), P(row_scale), P(tok), I(tok_stride as i32), P(hs), I(hb as i32), F(scale)])
	}
	pub fn spec_draft_pick(&self, drafts: u64, idx: u64, map: u64, j: usize, b: usize) -> Result<()> {
	self.be.launch("spec", "spec_draft_pick", [1, 1, 1], [b as u32, 1, 1], &[P(drafts), P(idx), P(map), I(j as i32)])
	}
	pub fn spec_draft_pos(&self, dpos: u64, len: u64, b: usize) -> Result<()> {
	self.be.launch("spec", "spec_draft_pos", [1, 1, 1], [b as u32, 1, 1], &[P(dpos), P(len)])
	}
	pub fn spec_h_gather(&self, h0: u64, hpool: u64, slot: u64, b: usize, hb: usize) -> Result<()> {
	self.be.launch("spec", "spec_h_gather", [b as u32, 1, 1], [256, 1, 1], &[P(h0), P(hpool), P(slot), I(hb as i32)])
	}
	pub fn spec_h_scatter(&self, hpool: u64, h0: u64, slot: u64, b: usize, hb: usize) -> Result<()> {
	self.be.launch("spec", "spec_h_scatter", [b as u32, 1, 1], [256, 1, 1], &[P(hpool), P(h0), P(slot), I(hb as i32)])
	}
	}