Class Native

java.lang.Object
smile.torch.Native
public final class Native extends Object
Internal helpers for the FFM (Foreign Function and Memory API) binding to the smile_torch native library. This class glues the idiomatic smile.deep.tensor API to the raw smile_torch_h downcall stubs generated by jextract.

It centralizes three concerns: converting native error sentinels into Java exceptions, freeing native handles when their Java wrappers become unreachable (via a shared Cleaner), and a handful of downcalls for functions that are not yet covered by the generated bindings.

  • Field Details

    • CLEANER

      public static final Cleaner CLEANER
      Shared cleaner that releases native handles owned by Java wrappers (tensors, indices, modules, …) once those wrappers are no longer reachable. It is a safety net; callers should still close resources deterministically via try-with-resources or an AutoScope where possible.

    • SCALING_TENSOR_WISE

      public static final int SCALING_TENSOR_WISE
      at::blas::ScalingType ordinals for scaledMmV2(Tensor, Tensor, Tensor, Tensor, int, int, int).
      See Also:

    • SCALING_ROW_WISE

      public static final int SCALING_ROW_WISE
      See Also:

    • SCALING_BLOCK_WISE_1X128

      public static final int SCALING_BLOCK_WISE_1X128
      See Also:

    • SCALING_BLOCK_WISE_128X128

      public static final int SCALING_BLOCK_WISE_128X128
      See Also:

  • Method Details

    • lastError

      public static String lastError()
      Returns the message describing the most recent native failure, or an empty string if none.
      Returns:
      the last native error message.

    • check

      public static MemorySegment check(MemorySegment handle)
      Verifies that a native call returned a non-null handle, throwing the native error message otherwise.
      Parameters:
      handle - the handle returned by a native call.
      Returns:
      the validated handle.

    • copy_

      public static void copy_(Tensor dst, Tensor src)
      Copies src into dst in-place.
      Parameters:
      dst - the destination tensor.
      src - the source tensor.

    • loadStateDict

      public static void loadStateDict(MemorySegment module, Map<String,Tensor> stateDict, boolean strict)
      Loads a flat state dict into a module by matching named_parameters() keys.
      Parameters:
      module - the native ST_Module handle.
      stateDict - map from fully-qualified parameter name to tensor.
      strict - when true, every module parameter must be present and every state-dict key must match a module parameter.

    • cudaMemGetInfo

      public static long[] cudaMemGetInfo(int deviceIndex)
      Queries free and total CUDA device memory via LibTorch's caching allocator (cudaMemGetInfo under a device guard).

      Requires a smile_torch binary built with USE_CUDA (CUDA Toolkit + CUDA LibTorch). A CPU-only build fails with a clear native error rather than a silent -1.

      Parameters:
      deviceIndex - the CUDA device index.
      Returns:
      long[2] of {freeBytes, totalBytes}.
      Throws:
      RuntimeException - if the query fails (e.g. CUDA unavailable).

    • cudaAllocatorStats

      public static long[] cudaAllocatorStats(int deviceIndex)
      Queries CUDACachingAllocator live vs reserved bytes for a device.

      allocated is storage owned by live tensors; reserved is the cudaMalloc footprint (includes free cached blocks that Device.emptyCache() can return to the driver).

      Parameters:
      deviceIndex - the CUDA device index.
      Returns:
      long[2] of {allocatedBytes, reservedBytes}.
      Throws:
      RuntimeException - if the query fails.

    • cudaComputeCapability

      public static int[] cudaComputeCapability(int deviceIndex)
      Returns CUDA compute capability {major, minor} for a device.
      Parameters:
      deviceIndex - CUDA device index.
      Returns:
      int[2] of {major, minor}.
      Throws:
      RuntimeException - if the query fails or the symbol is missing (rebuild libsmile_torch with CUDA).

    • scaledMm

      public static Tensor scaledMm(Tensor a, Tensor b, Tensor scaleA, Tensor scaleB, int outDtype)
      Scaled FP8 matrix multiply via LibTorch at::_scaled_mm (cuBLASLt).
      Parameters:
      a - FP8 activation [M,K].
      b - FP8 weight layout expected by _scaled_mm [N,K].
      scaleA - float scale for a.
      scaleB - float scale for b.
      outDtype - output ScalarType code (smile_torch_h.ST_DTYPE_BFLOAT16() etc.).
      Returns:
      output tensor, or null if the native path is unavailable.

    • scaledMmV2

      public static Tensor scaledMmV2(Tensor a, Tensor b, Tensor scaleA, Tensor scaleB, int recipeA, int recipeB, int outDtype)
      Block-scaled FP8 matrix multiply via LibTorch at::_scaled_mm_v2.
      Parameters:
      a - FP8 activation [M,K].
      b - FP8 weight [N,K].
      scaleA - float32 scales for a (BlockWise1x128 layout).
      scaleB - float32 scales for b (BlockWise128x128 layout).
      recipeA - ScalingType ordinal for a.
      recipeB - ScalingType ordinal for b.
      outDtype - output dtype code.
      Returns:
      output tensor, or null if unavailable.

    • scaledMmV2Available

      public static boolean scaledMmV2Available()
      Returns:
      true when scaledMmV2(Tensor, Tensor, Tensor, Tensor, int, int, int) is compiled into libsmile_torch.

    • fp8Quant1x128

      public static Tensor[] fp8Quant1x128(Tensor input)
      Dynamic BlockWise1x128 activation quantize on CUDA.
      Parameters:
      input - [M,K] with K % 128 == 0.
      Returns:
      {fp8 [M,K], scale [M, K/128]} or null if unavailable.

    • fp8Quant1x128Available

      public static boolean fp8Quant1x128Available()
      Returns:
      true when fp8Quant1x128(Tensor) is compiled into libsmile_torch.

    • marlinAvailable

      public static boolean marlinAvailable()
      Returns:
      true when Marlin is compiled into libsmile_torch.

    • marlinMul

      public static Tensor marlinMul(Tensor a, Tensor b, Tensor scales, Tensor workspace, int threadK)
      Marlin FP16×INT4 GEMM. Returns null when Marlin is not built.
      Parameters:
      a - FP16 activations [M,K].
      b - Marlin-packed INT4 weights.
      scales - group scales.
      workspace - scratch buffer (may be empty).
      threadK - Marlin thread_k (-1 = auto).

    • marlinFp8Available

      public static boolean marlinFp8Available()
      Returns:
      true when the MXFP8 Marlin kernels are compiled into libsmile_torch.

    • marlinMoeFp8Available

      public static boolean marlinMoeFp8Available()
      Returns:
      true when the MoE token-grouping and Marlin MoE GEMM kernels are available.

    • marlinFp8RepackWeight

      public static Tensor marlinFp8RepackWeight(Tensor weight)
      Repacks an MXFP8 weight into the Marlin tile layout.
      Parameters:
      weight - FP8 E4M3 (or uint8 holding E4M3 bytes) [N, K] on CUDA; N % 64 == 0, K % 128 == 0.
      Returns:
      int32 [K / 16, N * 4].

    • marlinFp8BlockScalesAvailable

      public static boolean marlinFp8BlockScalesAvailable()
      Returns:
      true when the block-FP8 scale repack (Qwen/DeepSeek checkpoints on Ampere) is compiled in.

    • marlinFp8RepackBlockScales

      public static Tensor marlinFp8RepackBlockScales(Tensor scaleInv)
      Repacks block-FP8 scales (W = fp8 * scale per 128x128 tile) for the group-128 Marlin kernel.
      Parameters:
      scaleInv - float32 [N/128, K/128] on CUDA.
      Returns:
      bfloat16 [K/128, N] in the Marlin layout, FP8 exponent bias folded in.

    • marlinFp8RepackScales

      public static Tensor marlinFp8RepackScales(Tensor scales)
      Permutes E8M0 scales into the Marlin layout.
      Parameters:
      scales - uint8 [N, K / 32].
      Returns:
      uint8 [K / 32, N].

    • marlinFp8Gemm

      public static Tensor marlinFp8Gemm(Tensor a, Tensor qweight, Tensor scales)
      MXFP8 weight-only GEMM: A @ dequant(W)^T.
      Parameters:
      a - bfloat16 activations [M, K] on CUDA.
      qweight - output of marlinFp8RepackWeight(Tensor).
      scales - output of marlinFp8RepackScales(Tensor).
      Returns:
      bfloat16 [M, N].

    • marlinFp8RepackWeightStack

      public static Tensor marlinFp8RepackWeightStack(Tensor weights)
      Stacked variant of marlinFp8RepackWeight(Tensor).
      Parameters:
      weights - FP8 [E, N, K].
      Returns:
      int32 [E, K / 16, N * 4].

    • marlinFp8RepackScalesStack

      public static Tensor marlinFp8RepackScalesStack(Tensor scales)
      Stacked variant of marlinFp8RepackScales(Tensor).
      Parameters:
      scales - uint8 [E, N, K / 32].
      Returns:
      uint8 [E, K / 32, N].

    • moeAlignBlockSize

      public static Tensor moeAlignBlockSize(Tensor topkIds, int expertStart, int numLocalExperts, int blockSize)
      Groups top-k assignments by expert on the device for the Marlin MoE kernel.
      Parameters:
      topkIds - expert ids [T, topK] (int64 or int32) on CUDA.
      expertStart - first expert owned by this rank.
      numLocalExperts - experts owned by this rank.
      blockSize - MoE block size (8, 16, 32, 48 or 64).
      Returns:
      one int32 tensor: sorted token ids, per-block expert ids, then the padded length.

    • marlinMoeFp8Gemm

      public static Tensor marlinMoeFp8Gemm(Tensor a, Tensor qweight, Tensor scales, Tensor align, Tensor topkWeights, int blockSize, int topK, boolean mulTopkWeights)
      MXFP8 Marlin MoE GEMM; see smile_marlin_fp8.h.
      Parameters:
      a - bfloat16 [M, K].
      qweight - int32 [E, K/16, N*4].
      scales - uint8 [E, K/32, N].
      align - result of moeAlignBlockSize(Tensor, int, int, int).
      topkWeights - float32, T * topK routing weights.
      blockSize - the block size passed to moeAlignBlockSize(Tensor, int, int, int).
      topK - top-k of the routing (1 for the second GEMM).
      mulTopkWeights - whether to scale each row by its routing weight.
      Returns:
      bfloat16 [M * topK, N].

    • fusedOpsAvailable

      public static boolean fusedOpsAvailable()
      Returns:
      true when the fused decode kernels are compiled in and not disabled by SMILE_FUSED_OPS=0.

    • gemmaRmsNorm

      public static Tensor gemmaRmsNorm(Tensor x, Tensor weight, double eps)
      Gemma RMSNorm x * rsqrt(mean(x^2) + eps) * (1 + w) in one launch.
      Returns:
      the normalized tensor, or null when the kernel is unavailable or the input is unsupported (not CUDA, unsupported dtype).

    • addGemmaRmsNorm

      public static Tensor[] addGemmaRmsNorm(Tensor x, Tensor r, Tensor weight, double eps)
      Residual add fused with Gemma RMSNorm: sum = x + r, y = gemmaRmsNorm(sum), one launch.
      Returns:
      {y, sum}, or null when unavailable or unsupported (nothing allocated).

    • linearWeight

      public static Tensor linearWeight(MemorySegment linear)
      The weight [out, in] of a smile_linear layer as an alias of the module parameter, or null.
      Parameters:
      linear - the ST_Linear handle.

    • gemv2

      public static Tensor[] gemv2(Tensor x, Tensor wa, Tensor wb)
      Two small bf16 projections x . wa^T and x . wb^T in one launch (at most 8 rows).
      Returns:
      {a, b} each [rows, N], or null when unsupported (nothing allocated).

    • sigmoidGateMul

      public static Tensor sigmoidGateMul(Tensor attn, Tensor gate)
      attn * sigmoid(gate) where gate is a strided [B,S,H,D] view; null when unsupported.

    • kvWrite

      public static boolean kvWrite(Tensor kLayer, Tensor vLayer, Tensor k, Tensor v, Tensor slots)
      Writes new K/V rows into one layer's caches at slots (int64) in one launch.
      Returns:
      true when written; false (nothing written) when the layout is unsupported.

    • siluMul

      public static Tensor siluMul(Tensor gate, Tensor up)
      silu(gate) * up in one launch, or null when unavailable or unsupported.

    • qkNormRope

      public static Tensor[] qkNormRope(Tensor q, Tensor k, Tensor qWeight, Tensor kWeight, Tensor cos, Tensor sin, int rotaryDim, double eps)
      Per-head Gemma RMSNorm of q and k then partial RoPE, one launch; see smile_minimax_ops.h.
      Returns:
      {q, k}, or null when unavailable or unsupported.

    • swigluOai

      public static Tensor swigluOai(Tensor gate, Tensor up, double alpha, double limit)
      The swigluoai activation in one launch.
      Returns:
      [T, I] activations, or null when unavailable or unsupported.

    • moeRoute

      public static Tensor[] moeRoute(Tensor logits, Tensor bias, int topK, double scale)
      MoE routing (sigmoid, bias-corrected top-k, normalise, scale) in one launch.
      Returns:
      {ids int32 [T, k], weights float32 [T, k]}, or null when unavailable or unsupported.

    • msaDecodeMask

      public static Tensor msaDecodeMask(Tensor q, Tensor cache, Tensor slots, Tensor positions, int length, int blockSize, int topk, int initBlocks, int localBlocks)
      MSA block selection for single-token decode; see smile_minimax_ops.h.
      Parameters:
      q - index queries [B, G, D].
      cache - one layer of the index-key cache [numSlots, 1, D], contiguous.
      slots - int64 [B * length] physical slots.
      positions - int32 [B] query positions on the device.
      Returns:
      additive mask [B, G, 1, length], or null when unavailable or unsupported.

    • msaDecodeSelect

      public static Tensor msaDecodeSelect(Tensor q, Tensor cache, Tensor slots, Tensor positions, int length, int blockSize, int topk, int initBlocks, int localBlocks, int seqlen)
      MSA block selection for single-token decode returning block ids; see smile_minimax_ops.h.
      Parameters:
      q - index queries [T, G, D], T = B * seqlen.
      seqlen - query tokens per row (1 for decode); token t is at positions[t / seqlen] + t % seqlen.
      Returns:
      int32 [T, G, topk] (-1 pads), or null when unavailable or unsupported.

    • sparseDecodeAttention

      public static Tensor sparseDecodeAttention(Tensor q, Tensor kLayer, Tensor vLayer, Tensor slots, Tensor positions, Tensor sel, int length, double scale, int seqlen)
      Block-sparse decode attention over the blocks chosen by msaDecodeSelect(Tensor, Tensor, Tensor, Tensor, int, int, int, int, int, int); see smile_minimax_ops.h.
      Parameters:
      q - bfloat16 [T, H, 128], T = B * seqlen.
      Returns:
      bfloat16 [T, H, 128], or null when unavailable or unsupported.

    • resizeBicubicAa

      public static Tensor resizeBicubicAa(Tensor x, int outH, int outW)
      Bicubic resize with antialiasing (torchvision resize(antialias=True)).
      Parameters:
      x - [N, C, H, W], uint8 or floating point, on any device.
      outH - target height.
      outW - target width.
      Returns:
      [N, C, outH, outW] of the same dtype.
      Throws:
      IllegalStateException - if the loaded native library predates this function.

    • mxfp8ToBlockFp8Available

      public static boolean mxfp8ToBlockFp8Available()
      Returns:
      true when the MXFP8 to block-FP8 converter is compiled into libsmile_torch.

    • mxfp8ToBlockFp8

      public static Tensor[] mxfp8ToBlockFp8(Tensor weight, Tensor scale)
      Re-expresses MXFP8 weights (1x32 E8M0 scales) as block FP8 (one float32 scale per 128x128 tile).
      Parameters:
      weight - FP8 E4M3 [N, K] or [E, N, K] on CUDA, N and K multiples of 128.
      scale - uint8 E8M0 scales [.., N, K / 32].
      Returns:
      {fp8, scaleInv} where scaleInv is float32 [.., N/128, K/128] and W = fp8 * scaleInv per tile.

    • moduleSetRequiresGrad

      public static void moduleSetRequiresGrad(MemorySegment module, boolean requiresGrad)
      Sets requires_grad on every parameter of module (recursive).
      Parameters:
      module - native ST_Module handle.
      requiresGrad - true to enable autograd on parameters.

    • unregisterModule

      public static void unregisterModule(MemorySegment module, String name)
      Unregisters a child module by name, dropping the parent's shared_ptr. Used when replacing dense LinearLayer shells with quantized ops so GPU storage is released before KV pool sizing.
      Parameters:
      module - parent ST_Module.
      name - registered child name (e.g. "wq").

    • nbytes

      public static long nbytes(Tensor tensor)
      Returns the storage size of tensor in bytes.
      Parameters:
      tensor - the tensor.
      Returns:
      tensor.nbytes().

    • dataPtr

      public static MemorySegment dataPtr(Tensor tensor)
      Returns a writable MemorySegment covering the contiguous storage of tensor (nbytes bytes). The tensor must be contiguous and remain reachable for the lifetime of the segment.
      Parameters:
      tensor - the contiguous tensor.
      Returns:
      the storage segment.

    • tpAllReduceSum

      public static void tpAllReduceSum(Tensor[] tensors)
      In-place sum all-reduce across tensors that already reside on distinct devices.
      Parameters:
      tensors - one tensor per TP rank (same shape/dtype).

    • tpBroadcast

      public static void tpBroadcast(Tensor[] tensors, int root)
      Broadcasts tensors[root] onto every other entry.

    • ncclCommCreate

      public static MemorySegment ncclCommCreate(int[] deviceIndices)
      Creates a local NCCL communicator for deviceIndices, or returns MemorySegment.NULL when NCCL is unavailable.

    • ncclUniqueId

      public static byte[] ncclUniqueId()
      A fresh NCCL unique id (128 bytes) for a communicator that spans processes.
      Returns:
      the id.
      Throws:
      IllegalStateException - when NCCL is unavailable.

    • ncclCommCreateRanks

      public static MemorySegment ncclCommCreateRanks(int[] deviceIndices, int worldSize, int rankOffset, byte[] uniqueId)
      Creates this process's ranks of a communicator spanning worldSize ranks: deviceIndices.length local ranks starting at world rank rankOffset. Every process must call this at the same time with the same uniqueId; it returns once all ranks have joined. The collectives take local rank indices.
      Returns:
      the communicator handle.
      Throws:
      IllegalStateException - when it cannot be created.

    • ncclCommFree

      public static void ncclCommFree(MemorySegment comm)
      Destroys an NCCL communicator from ncclCommCreate(int[]).

    • ncclAllReduceSum

      public static void ncclAllReduceSum(MemorySegment comm, int rank, Tensor local)
      In-place NCCL sum all-reduce on local for rank. Every TP rank must call concurrently.

    • ncclBroadcast

      public static void ncclBroadcast(MemorySegment comm, int rank, int root, Tensor local)
      NCCL broadcast of local from root. Every rank must call concurrently with its buffer.

    • recurrentGatedDeltaRule

      public static Tensor recurrentGatedDeltaRule(Tensor query, Tensor key, Tensor value, Tensor g, Tensor beta, Tensor state, boolean qkL2norm)
      Fused recurrent gated delta rule. Mutates float state in place. Returns output [B,S,H,Dv] (caller owns).
      Returns:
      output tensor, or null if the native path is unavailable.

    • hasGatedDeltaVerifyWindow

      public static boolean hasGatedDeltaVerifyWindow()
      Whether gatedDeltaVerifyWindow(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, int, int, int, int, boolean, Tensor[], Tensor[]) is available in the loaded native library.
      Returns:
      true when the symbol resolved.

    • gatedDeltaVerifyWindow

      public static Tensor gatedDeltaVerifyWindow(Tensor hidden, Tensor convState, Tensor convWeight, Tensor g, Tensor beta, Tensor recState, int numKHeads, int numVHeads, int headKDim, int headVDim, boolean qkL2norm, Tensor[] convCkpts, Tensor[] recCkpts)
      Fused MTP verify window for one DeltaNet layer: conv + SiLU + QKV split in one kernel (rolls convState, emits per-position conv checkpoints) followed by the fused recurrent rule (updates recState, emits per-position state checkpoints). CUDA only.
      Parameters:
      hidden - pre-conv QKV projection [B, S, C] (contiguous).
      g - decay gate [B, S, Hv].
      beta - input gate [B, S, Hv].
      Returns:
      core output [B, S, Hv, Dv] in hidden's dtype, or null when unsupported (nothing modified) so the caller can use another path.

    • hasGatedDeltaVerifyWindowFused

      public static boolean hasGatedDeltaVerifyWindowFused()
      Whether gatedDeltaVerifyWindowFused(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, int, int, int, int, Tensor[], Tensor[]) is available.

    • gatedDeltaVerifyWindowFused

      public static Tensor gatedDeltaVerifyWindowFused(Tensor hidden, Tensor convState, Tensor convWeight, Tensor a, Tensor b, Tensor aLogF, Tensor dtBiasF, Tensor recState, int numKHeads, int numVHeads, int headKDim, int headVDim, Tensor[] convCkpts, Tensor[] recCkpts)
      gatedDeltaVerifyWindow(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, int, int, int, int, boolean, Tensor[], Tensor[]) with the gates computed inside the recurrent kernel from the raw a/b projections [B, S, Hv] (bf16) and float aLogF/dtBiasF; convWeight should be a bf16 [C, K] contiguous copy so no per-call conversion is needed.
      Returns:
      core [B, S, Hv, Dv], or null when unsupported (nothing modified).

    • hasRecurrentGatedDeltaRuleCkpt

      public static boolean hasRecurrentGatedDeltaRuleCkpt()
      Whether recurrentGatedDeltaRuleCkpt(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, boolean, Tensor[]) is available in the loaded native library.
      Returns:
      true when the symbol resolved.

    • recurrentGatedDeltaRuleCkpt

      public static Tensor recurrentGatedDeltaRuleCkpt(Tensor query, Tensor key, Tensor value, Tensor g, Tensor beta, Tensor state, boolean qkL2norm, Tensor[] ckpts)
      Fused recurrent gated delta rule over a whole window that also writes the state after step t into ckpts[t] (CUDA only). Mutates state to the end-of-window state.
      Parameters:
      ckpts - float, contiguous per-step state sinks (at least as large as state).
      Returns:
      output [B,S,H,Dv] (caller owns), or null when unavailable (older library or non-fused fallback) so the caller can use a per-position loop.

    • softplus

      public static Tensor softplus(Tensor x)
      Native torch::softplus. Returns null when the symbol is missing (older libsmile_torch).

    • gatedDeltaComputeG

      public static Tensor gatedDeltaComputeG(Tensor a, Tensor aLog, Tensor dtBias)
      Fused g = -exp(A_log) * softplus(a + dt_bias) (float). Returns null when unavailable.

    • gatedDeltaComputeGates

      public static Tensor[] gatedDeltaComputeGates(Tensor a, Tensor b, Tensor aLog, Tensor dtBias)
      Fused beta = sigmoid(b) and decay gate g. Returns [g, beta] or null when unavailable.

    • repeatKvHeads

      public static Tensor repeatKvHeads(Tensor x, int rep)
      Repeat heads on dim 2 ([B,S,H,D] → [B,S,H*rep,D]). Returns null when unavailable.

    • rmsNorm

      public static Tensor rmsNorm(Tensor x, Tensor weight, double eps)
      Qwen RMSNorm rms_norm(x)*(1+weight). Returns null when unavailable.

    • rmsNormGated

      public static Tensor rmsNormGated(Tensor x, Tensor gate, Tensor weight, double eps)
      Fused gated RMSNorm. Returns null when unavailable.

    • mulSigmoid

      public static Tensor mulSigmoid(Tensor x, Tensor gate)
      x * sigmoid(gate). Returns null when unavailable.

    • causalConv1dUpdate

      public static Tensor causalConv1dUpdate(Tensor hidden, Tensor convState, Tensor weight)
      Fused causal depthwise conv1d update for decode (S==1): writes SiLU output and rolls convState in place. Returns null when unavailable.

    • causalConv1dUpdateSplitQkv

      public static Tensor[] causalConv1dUpdateSplitQkv(Tensor hidden, Tensor convState, Tensor weight, int numKHeads, int numVHeads, int headKDim, int headVDim)
      Fused decode conv + QKV split + head repeat. Returns [query, key, value] with shapes [B,1,numVHeads,headKDim] / [B,1,numVHeads,headVDim], or null when unavailable.

    • hasGatedDeltaDecodeFused

      public static boolean hasGatedDeltaDecodeFused()
      Whole Gated DeltaNet decode core (S == 1, bf16): conv update + QKV split + gates + recurrent rule in two launches. Updates convState and recState in place.
      Returns:
      the core [B * numVHeads, headVDim], or null (nothing modified) when unsupported.

    • gatedDeltaDecodeFused

      public static Tensor gatedDeltaDecodeFused(Tensor hidden, Tensor convState, Tensor convWeightF, Tensor a, Tensor b, Tensor aLogF, Tensor dtBiasF, Tensor recState, int numKHeads, int numVHeads, int headKDim, int headVDim)

    • applyRotaryPosEmb

      public static Tensor[] applyRotaryPosEmb(Tensor xq, Tensor xk, Tensor cos, Tensor sin, int rotaryDim)
      HF apply_rotary_pos_emb / rotate_half on query and key. Returns [q, k] or null when unavailable.

    • flashInferSetAotDir

      public static void flashInferSetAotDir(String path)
      Sets the FlashInfer AOT / jit-cache directory for the native loader.
      Parameters:
      path - absolute directory, or null/empty to clear.

    • flashInferAvailable

      public static boolean flashInferAvailable()
      Returns:
      true when FlashInfer paged attention is compiled into libsmile_torch.

    • flashInferWorkspaceCreate

      public static MemorySegment flashInferWorkspaceCreate(int deviceIndex, long workspaceBytes)
      Creates a FlashInfer workspace on deviceIndex.
      Returns:
      native handle, or null when unavailable.

    • flashInferWorkspaceFree

      public static void flashInferWorkspaceFree(MemorySegment handle)
      Frees a FlashInfer workspace handle.

    • flashInferWorkspaceInvalidateRuntimeCache

      public static void flashInferWorkspaceInvalidateRuntimeCache(MemorySegment handle)
      Clears cached FlashInfer decode plans and prefill gather slots.

    • flashInferWorkspaceInvalidatePrefillRuntimeCache

      public static void flashInferWorkspaceInvalidatePrefillRuntimeCache(MemorySegment handle)
      Clears only prefill gather caches on the workspace.

    • flashInferWorkspaceInvalidateVerifyRuntimeCache

      public static void flashInferWorkspaceInvalidateVerifyRuntimeCache(MemorySegment handle)
      Clears only the cached verify-graph (SMILE_VERIFY_CUDA_GRAPH) plan. Independent of flashInferWorkspaceInvalidateRuntimeCache(MemorySegment) / flashInferWorkspaceInvalidatePrefillRuntimeCache(MemorySegment) — decode's routine cohort/CSR-rebuild invalidation must not clear a still-valid verify plan, and vice versa.

    • flashInferAttentionPagedRaw

      public static Tensor flashInferAttentionPagedRaw(Tensor query, Tensor kCache, Tensor vCache, Tensor kvIndptr, Tensor kvIndices, Tensor kvLastPageLen, int pageSize, int numKvHeads, int headDim, int cacheLen, double scale, boolean isCausal, Tensor mask, MemorySegment workspace)
      Raw-tensor paged attention (existing eager S==1/S>1 dispatch), bypassing AttentionContext/KvCachePool.

      Exists so a Stage 1 (SMILE_VERIFY_CUDA_GRAPH) smoke test can compare flashInferAttentionVerifyCapturable(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, int, int, int, int, double, float, float, boolean, MemorySegment)'s output against this existing, proven path on identical raw tensors, without needing to stand up the full production KV-pool machinery. flashInferAttention(Tensor, Tensor, AttentionContext) is the production entry point built on the same underlying call; this is a thin sibling at the same tensor-level boundary.

      Parameters:
      query - query [B, Hq, S, D]
      kCache - key cache [numSlots, Hkv, D] (one layer slice)
      vCache - value cache [numSlots, Hkv, D]
      kvIndptr - int32 [B+1]
      kvIndices - int32 [numPages]
      kvLastPageLen - int32 [B]
      pageSize - tokens per page
      numKvHeads - Hkv
      headDim - D
      cacheLen - total sequence length (for CSR validation; 0 to skip)
      scale - attention scale (<= 0 → 1/sqrt(D))
      isCausal - whether to apply causal masking; ignored (SDPA requires the caller not set both) when mask != null
      mask - optional explicit additive mask, e.g. [S, cacheLen] (broadcastable); pass non-null for "continuation window" causal semantics (bottom-right-aligned: column j valid for row i iff j <= i + (cacheLen - S)) — plain isCausal=true with no mask uses PyTorch SDPA's own is_causal, which is top-left-aligned for S != cacheLen and therefore wrong for this use case. May be null.
      workspace - FlashInfer workspace handle
      Returns:
      output [B, Hq, S, D]

    • flashInferAttention

      public static Tensor flashInferAttention(Tensor query, Tensor mask, AttentionContext ctx)
      Runs paged attention using AttentionContext.
      Parameters:
      query - query [B, Hq, S, D]
      mask - optional additive attention mask (same as torch_native); may be null
      ctx - paged context with pool + CSR + workspace
      Returns:
      output [B, Hq, S, D]

    • flashInferAttentionVerifyCapturable

      public static Tensor flashInferAttentionVerifyCapturable(Tensor query, Tensor kCache, Tensor vCache, Tensor qoIndptr, Tensor kvIndptr, Tensor kvIndices, Tensor kvLastPageLen, int pageSize, int numKvHeads, int headDim, int qoLen, double scale, float kScale, float vScale, boolean isCausal, MemorySegment workspace)
      Stage 1 (SMILE_VERIFY_CUDA_GRAPH): graph-capturable multi-token (S>1) causal paged attention for MTP window verify (see smile_flashinfer_paged_attention_verify_cuda). Isolated sibling of flashInferAttention(Tensor, Tensor, AttentionContext) — not wired into the live verify path yet; this exists so Stage 1's eager-mode correctness can be validated directly against flashInferAttention(Tensor, Tensor, AttentionContext)'s existing SDPA path before any later stage touches Qwen/QwenModel.

      No additive mask parameter (FlashInfer's masking here is compile-time MaskMode, unlike the additive-mask SDPA path).

      Parameters:
      query - query [B, Hq, S, D]
      kCache - key cache [numSlots, Hkv, D] (one layer slice)
      vCache - value cache [numSlots, Hkv, D]
      qoIndptr - int32 [B+1], values {0,S,2S,...}
      kvIndptr - int32 [B+1]
      kvIndices - int32 [numPages]
      kvLastPageLen - int32 [B]
      pageSize - tokens per page
      numKvHeads - Hkv
      headDim - D
      qoLen - S (must equal query's 3rd dim)
      scale - attention scale (<= 0 → 1/sqrt(D))
      kScale - FP8 KV key dequant scale (1.0 when KV is bf16/fp16)
      vScale - FP8 KV value dequant scale (1.0 when KV is bf16/fp16)
      isCausal - diagnostic-only toggle (MaskMode::kCausal vs kNone); production callers always pass true
      workspace - FlashInfer workspace handle
      Returns:
      output [B, Hq, S, D]

    • hasFlashInferAttentionVerify

      public static boolean hasFlashInferAttentionVerify()
      Whether the FlashInfer paged-prefill verify entry point is linked.
      Returns:
      true when flashInferAttentionVerifyGraph(Tensor, AttentionContext, Tensor) can be called.

    • flashInferAttentionVerifyGraph

      public static Tensor flashInferAttentionVerifyGraph(Tensor query, AttentionContext ctx, Tensor qoIndptr)
      Production entry point for the graph-capturable MTP window-verify kernel, built on AttentionContext (KV pool layer slice + CSR + workspace) — the verify-graph sibling of flashInferAttention(Tensor, Tensor, AttentionContext), targeting smile_flashinfer_paged_attention_verify instead of smile_flashinfer_paged_attention. ctx.isCausal() must be true (this kernel has no additive-mask parameter; causal masking is the compile-time, bottom-right-aligned MaskMode::kCausal validated in Stage 1/2).
      Parameters:
      query - query [B, Hq, S, D]
      ctx - paged context with pool + CSR + workspace (seqLen is the query window length)
      qoIndptr - int32 [B+1], values {0,S,2S,...} (KvCachePool.verifyQoIndptrBuf(int, int))
      Returns:
      output [B, Hq, S, D]

    • flashInferAttentionVerifyEager

      public static Tensor flashInferAttentionVerifyEager(Tensor query, AttentionContext ctx, Tensor qoIndptr)
      Verify attention for calls that are not part of a CUDA graph bucket (eager verify windows, batched/ragged verify). Plans into a scratch workspace disjoint from flashInferAttentionVerifyGraph(Tensor, AttentionContext, Tensor)'s, so a captured graph's baked plan arrays can never be overwritten by an eager plan of a different shape. Falls back to the graph entry point on an older native library.

    • flashInferRaggedAttention

      public static Tensor flashInferRaggedAttention(Tensor query, Tensor key, Tensor value, Tensor mask, AttentionContext ctx)
      Runs ragged contiguous self-attention (BatchPrefillWithRaggedKVCache or SDPA fallback).
      Parameters:
      query - [N, H, D] NHD
      key - [N, H, D]
      value - [N, H, D]
      mask - optional additive mask; may be null
      ctx - ragged context with AttentionContext.raggedIndptr()
      Returns:
      output [N, H, D]

    • cudaGraphAvailable

      public static boolean cudaGraphAvailable()
      Returns:
      true when smile_cuda_graph_create is exported.

    • cudaGraphCreate

      public static MemorySegment cudaGraphCreate()
      Returns:
      opaque graph handle, or null when unavailable.

    • cudaGraphDestroy

      public static void cudaGraphDestroy(MemorySegment handle)
      Parameters:
      handle - graph from cudaGraphCreate().

    • cudaGraphCaptureBegin

      public static void cudaGraphCaptureBegin(MemorySegment handle, int deviceIndex)
      Parameters:
      handle - graph handle. @param deviceIndex CUDA device ordinal.

    • cudaGraphCaptureEnd

      public static void cudaGraphCaptureEnd(MemorySegment handle)
      Parameters:
      handle - graph handle after cudaGraphCaptureBegin(MemorySegment, int).

    • cudaGraphReplay

      public static void cudaGraphReplay(MemorySegment handle)
      Parameters:
      handle - instantiated graph handle.

    • cudaGraphIsReady

      public static boolean cudaGraphIsReady(MemorySegment handle)
      Parameters:
      handle - graph handle.