FlashInferWorkspace

java.lang.Object

smile.llm.attention.FlashInferWorkspace

All Implemented Interfaces:: AutoCloseable

public final class FlashInferWorkspace extends Object implements AutoCloseable

Opaque FlashInfer workspace / plan handle (one per device / TP rank).

Method Summary

Modifier and Type

Method

Description

void

close()

static FlashInferWorkspace

create(int deviceIndex, long workspaceBytes)

Allocates a workspace on the given CUDA device.

MemorySegment

handle()

Returns the native workspace handle for downcalls.

void

invalidatePrefillRuntimeCache()

Clears prefill gather caches only; decode-plan cache may remain when still valid.

void

invalidateRuntimeCache()

Invalidates step-scoped decode-plan and prefill-gather caches on this workspace.

void

invalidateVerifyRuntimeCache()

Clears the cached verify-graph (SMILE_VERIFY_CUDA_GRAPH) prefill plan only.

Methods inherited from class Object
clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Method Details
- create
  
  public static FlashInferWorkspace create(int deviceIndex, long workspaceBytes)
  
  Allocates a workspace on the given CUDA device.
  
  Parameters:
  
  deviceIndex - CUDA device ordinal.
  
  workspaceBytes - scratch buffer size hint (0 = kernel default).
  
  Returns:
  
  workspace, or null when FlashInfer is not compiled in.
- handle
  
  public MemorySegment handle()
  
  Returns the native workspace handle for downcalls.
  
  Returns:
  
  native handle for downcalls.
- close
  
  public void close()
  
  Specified by:
  
  close in interface AutoCloseable
- invalidateRuntimeCache
  
  public void invalidateRuntimeCache()
  
  Invalidates step-scoped decode-plan and prefill-gather caches on this workspace. Safe to call when FlashInfer is unavailable (no-op).
- invalidateVerifyRuntimeCache
  
  public void invalidateVerifyRuntimeCache()
  
  Clears the cached verify-graph (SMILE_VERIFY_CUDA_GRAPH) prefill plan only. Must be called whenever the KV CSR tensors are rebuilt: the native plan cache's pointer-identity fast path skips comparing kv_indptr contents, and the caching allocator routinely hands the rebuilt tensor the freed tensor's address, so without this a stale plan (old page count) is silently reused for the new, larger bucket.
- invalidatePrefillRuntimeCache
  
  public void invalidatePrefillRuntimeCache()
  
  Clears prefill gather caches only; decode-plan cache may remain when still valid.