Class DecodeCudaGraph

java.lang.Object
smile.llm.engine.DecodeCudaGraph
public final class DecodeCudaGraph extends Object
Phase 2c: optional CUDA graph capture for batch-1 decode forward.

Unlike SMILE_FLASHINFER_DECODE_CUDA_GRAPH (reverted — FlashInfer plan flag only), this captures the full decode forward via at::cuda::CUDAGraph.

Enable with environment variable SMILE_DECODE_CUDA_GRAPH=1.

Phase 2d (multi-batch): graphs are bucketed by (batch, numPages) for uniform decode steps: batch must be a power of two in [1, maxBatch], and every row must share the same KV cache length. RoPE, token, logits, and KV index buffers are sized to the live batch.

  • Method Summary

    Modifier and Type
    Method
    Description
    static boolean
    canGraphDecode(int[] cachePositions)
    Returns whether this decode step is eligible for CUDA graph capture/replay.
    static void
    disableCapture(String reason)
    Permanently disable decode CUDA graphs for this process (after capture failure).
    static void
    disablePreCapture(String reason)
    Permanently disable next-bucket prefetch after a capture / replay failure.
    static boolean
    enabled()
    Returns whether decode CUDA graphs are enabled for this process.
    static boolean
    hasPrefetchHeadroom(long freeBytes)
    Returns whether free device memory is sufficient for prefetch capture.
    static void
    logPrefetchSkippedLowMemory(int tpRank, long freeBytes)
    Logs once when prefetch is skipped due to low free memory.
    static void
    markPersistentLogits(boolean persistent)
    Marks that the current decode step returned logits backed by a captured CUDA graph output buffer (must not be closed by the caller).
    static int
    maxBatch()
    Returns the maximum batch size eligible for decode CUDA graphs.
    static boolean
    persistentLogits()
    Returns whether decode logits outlive the logits-row copy step.
    static boolean
    preCaptureEnabled()
    Returns whether next-bucket prefetch capture is enabled.
    static int
    prefetchLeadSteps()
    Returns decode steps before a KV page boundary used to spread prefetch work.
    static long
    prefetchMinFreeBytes()
    Returns the configured prefetch free-memory floor.
    static boolean
    supportsBatch(int batch)
    Returns whether batch is a supported graph bucket size.
    static int
    warmupSteps()
    Returns the number of eager warmup decode steps before graph capture.

    Methods inherited from class Object

    clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

  • Method Details

    • enabled

      public static boolean enabled()
      Returns whether decode CUDA graphs are enabled for this process.
      Returns:
      true when env is set, native API is linked, and capture is not disabled.

    • maxBatch

      public static int maxBatch()
      Returns the maximum batch size eligible for decode CUDA graphs.
      Returns:
      maximum power-of-two batch size for graph buckets.

    • supportsBatch

      public static boolean supportsBatch(int batch)
      Returns whether batch is a supported graph bucket size.
      Parameters:
      batch - decode batch size.
      Returns:
      true when batch is a supported graph bucket size.

    • canGraphDecode

      public static boolean canGraphDecode(int[] cachePositions)
      Returns whether this decode step is eligible for CUDA graph capture/replay.
      Parameters:
      cachePositions - per-row KV write positions.
      Returns:
      true when every row shares the same KV write position and the batch size is graph-eligible.

    • disableCapture

      public static void disableCapture(String reason)
      Permanently disable decode CUDA graphs for this process (after capture failure).
      Parameters:
      reason - human-readable failure reason (for callers; not logged here).

    • warmupSteps

      public static int warmupSteps()
      Returns the number of eager warmup decode steps before graph capture.
      Returns:
      warmup steps per KV-page bucket.

    • preCaptureEnabled

      public static boolean preCaptureEnabled()
      Returns whether next-bucket prefetch capture is enabled.
      Returns:
      true when the next numPages bucket may be captured ahead of the page boundary (SMILE_DECODE_CUDA_GRAPH_PRE_CAPTURE=1).

    • disablePreCapture

      public static void disablePreCapture(String reason)
      Permanently disable next-bucket prefetch after a capture / replay failure.
      Parameters:
      reason - human-readable failure reason (for callers; not logged here).

    • hasPrefetchHeadroom

      public static boolean hasPrefetchHeadroom(long freeBytes)
      Returns whether free device memory is sufficient for prefetch capture.
      Parameters:
      freeBytes - free device memory in bytes; negative skips the check.
      Returns:
      true when the device has enough free memory for a safe next-bucket prefetch capture attempt.

    • prefetchMinFreeBytes

      public static long prefetchMinFreeBytes()
      Returns the configured prefetch free-memory floor.
      Returns:
      configured prefetch free-memory floor in bytes.

    • logPrefetchSkippedLowMemory

      public static void logPrefetchSkippedLowMemory(int tpRank, long freeBytes)
      Logs once when prefetch is skipped due to low free memory.
      Parameters:
      tpRank - tensor-parallel rank for logging.
      freeBytes - free device memory in bytes.

    • prefetchLeadSteps

      public static int prefetchLeadSteps()
      Returns decode steps before a KV page boundary used to spread prefetch work.
      Returns:
      prefetch lead steps (warmupSteps() + 1).

    • markPersistentLogits

      public static void markPersistentLogits(boolean persistent)
      Marks that the current decode step returned logits backed by a captured CUDA graph output buffer (must not be closed by the caller).
      Parameters:
      persistent - true when logits outlive the logits-row copy step.

    • persistentLogits

      public static boolean persistentLogits()
      Returns whether decode logits outlive the logits-row copy step.
      Returns:
      true when decode logits outlive the logits-row copy step.