Docs / jax4j-core / com.nablatensor.jax4j.tensor.spi
interface
QuantizedOps
Optional capability: hold quantized weights and multiply them by an F32 vector without ever materialising the dequantized matrix.
Discover with instanceof; never assume a ComputeBackend implements this. Added so tensorjit-llm can run fused dequant-GEMV kernels without adding quant-specific methods to the stable ComputeBackend SPI (see JAX4J_LLM_PLAN.md §4).
Methods
Set<QuantType> supportedQuantTypes()
DeviceBuffer uploadQuantized(MemorySegment blocks, QuantType type, Shape logicalShape, Device device)
Uploads an opaque quantized blob. logicalShape is the F32 shape it decodes to.
DeviceBuffer quantMatvec(DeviceBuffer weights, DeviceBuffer x)
y(rows) = W(rows, cols) . x(cols), dequantizing on the fly.
DeviceBuffer quantMatvecAdd(DeviceBuffer weights, DeviceBuffer x, DeviceBuffer add)
y(rows) = W(rows, cols) . x(cols) + add(rows) -- the fused residual.
DeviceBuffer quantGatherRow(DeviceBuffer weights, int row)
One logical row of a quantized matrix, as F32 -- the embedding lookup.