Struct Glm5NextKdaWorkspace 

pub struct Glm5NextKdaWorkspace {

    pub qkv_parts: DevicePtr,
    pub qkv_proj: DevicePtr,
    pub conv_out: DevicePtr,
    pub q_f32: DevicePtr,
    pub k_f32: DevicePtr,
    pub v_f32: DevicePtr,
    pub gate: DevicePtr,
    pub beta: DevicePtr,
    pub core: DevicePtr,
    pub g_raw: DevicePtr,
    pub out_gate: DevicePtr,
    pub o_norm_out: DevicePtr,
    pub final_out: DevicePtr,
    /* private fields */
}

Scratch owned by the forward path. Sized once for max_tokens and reused across layers — the whole KDA family shares one workspace because every block has identical geometry.

The intermediate buffers are pub on purpose: the numeric oracle compares stage by stage, and a residual quoted only at the layer output cannot separate a kernel bug from a rounding floor.

Fields

qkv_parts: DevicePtr

[3, T, qkv] BF16 — the three projections as dense_gemm_bf16 writes them.

qkv_proj: DevicePtr

[T, conv_dim] BF16 — q|k|v per token, pre-conv.

conv_out: DevicePtr

[T, conv_dim] BF16 — post conv + SiLU. On the decode path L2 is already fused in.

q_f32: DevicePtr

[T_pad, qkv] FP32 — post-L2 q, post-L2 k, raw v. Prefill only.

k_f32: DevicePtr

v_f32: DevicePtr

gate: DevicePtr

[T_pad, heads, head_dim] FP32 — bounded log-decay from kda_gate.

beta: DevicePtr

[T_pad, heads] FP32 — already sigmoided.

core: DevicePtr

[T_pad, heads, head_dim] FP32 — KDA core output, pre-norm.

g_raw: DevicePtr

[T, qkv] BF16 — f_b(f_a(h)), the forget-gate projection kda_gate consumes. Kept separate from out_gate: same shape, same kind of low-rank pair, aliasing is silent.

out_gate: DevicePtr

[T, qkv] BF16 — g_b(g_a(h)), the low-rank output gate.

o_norm_out: DevicePtr

[T, qkv] BF16 — after the sigmoid-gated RMSNorm.

final_out: DevicePtr

[T, hidden] BF16 — the block output.

Implementations

impl Glm5NextKdaWorkspace

pub fn new( gpu: &dyn GpuBackend, cfg: &Glm5NextKdaConfig, max_tokens: usize, ) -> Result<Self>

pub fn max_tokens(&self) -> usize

pub fn t_pad(&self) -> usize