How to Enable and Interpret Auto-Dispatch Logs for FlashKDA and FLA Integration

Enable Python’s INFO-level logging to view [FLA Backend] messages that indicate whether FlashKDA or the Triton fallback is handling your chunk_kda operations.

FlashKDA integrates with flash-linear-attention (FLA) as a high-performance CUDA backend for the chunk_kda operation. When you enable and interpret the auto-dispatch logs, you can verify that FlashKDA is active or diagnose exactly why the system fell back to Triton kernels.

Enabling Auto-Dispatch Logging

FlashKDA relies on FLA’s dispatch mechanism, which emits diagnostic messages through Python’s standard logging module. To capture these messages, configure your script to display INFO-level logs before importing FLA:

import logging
logging.basicConfig(level=logging.INFO)

This configuration triggers the dispatch logic in the FLA repository to emit one of two messages for every chunk_kda call: either confirming FlashKDA selection or explaining its rejection.

Interpreting the Log Output

Once enabled, the logs reveal the backend selection outcome immediately before kernel execution.

Successful Backend Selection (Hit)

When FlashKDA is successfully selected, the log contains:

[FLA Backend] kda.chunk_kda -> flashkda

This message confirms that the auto-dispatch succeeded and that the CUDA implementation in csrc/flash_kda.cpp will handle the computation via the flash_kda.fwd entry point exposed in flash_kda/__init__.py (lines 5-6).

Dispatch Rejection (Miss)

If FlashKDA cannot be used, the log shows a rejection reason:

[FLA Backend] kda.chunk_kda rejected: <reason>

Common rejection reasons include unsupported GPU architecture (SM90+ required), dtype mismatches, or missing environment variables. The text after rejected: explains exactly why FlashKDA was skipped, allowing you to adjust your configuration or hardware setup.

Disabling FlashKDA for Debugging

To force FLA to use the Triton implementation and bypass FlashKDA entirely, set the environment variable before importing the library. This is documented in README.md (lines 62-64):

import os
os.environ["FLA_FLASH_KDA"] = "0"

This is useful for performance comparisons or when debugging kernel-specific issues by ensuring the dispatch logic selects the reference implementation.

Practical Code Examples

Minimal Script with Logging

The following example demonstrates how to enable logs and invoke the operation. The output will indicate whether FlashKDA or Triton is used:

import logging
import torch
from fla.ops.kda import chunk_kda

# Enable INFO-level logs

logging.basicConfig(level=logging.INFO)

# Example tensors (CUDA, bf16, K=V=128)

B, T, H, K = 1, 128, 8, 128
q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
k = torch.randn_like(q)
v = torch.randn_like(q)
g = torch.randn_like(q)
beta = torch.randn(B, T, H, dtype=torch.bfloat16, device='cuda')
scale = 1.0
A_log = torch.randn(H, dtype=torch.float32, device='cuda')
dt_bias = torch.randn(H, K, dtype=torch.float32, device='cuda')
lower_bound = -5.0

# Dispatch occurs here; check logs for backend selection

out, _ = chunk_kda(
    q, k, v, g, beta,
    scale=scale,
    initial_state=None,
    final_state=None,
    A_log=A_log,
    dt_bias=dt_bias,
    lower_bound=lower_bound,
    cu_seqlens=None,
)

Forcing Triton Fallback

To verify behavior without FlashKDA:

import os
os.environ["FLA_FLASH_KDA"] = "0"

import torch
from fla.ops.kda import chunk_kda

# ... remainder of the script

Key Source Files

• README.md (lines 62-64): Documents the FLA_FLASH_KDA environment variable and the debug-dispatch logging hook.
• flash_kda/__init__.py (lines 5-6): Exposes the Python entry point flash_kda.fwd that FLA calls after successful dispatch.
• csrc/flash_kda.cpp (lines 1-33): Implements the CUDA kernel and workspace logic utilized when auto-dispatch succeeds.
• flash-linear-attention (external): Contains the dispatch wrapper that emits the [FLA Backend] log messages and performs the backend selection logic.

Summary

• Enable logs by setting logging.basicConfig(level=logging.INFO) to capture FLA’s backend selection messages.
• Interpret hits via [FLA Backend] kda.chunk_kda -> flashkda confirming CUDA kernel usage.
• Diagnose misses by reading the rejected: reason to identify hardware or configuration issues.
• Disable FlashKDA by setting FLA_FLASH_KDA=0 to force Triton execution for comparison testing.
• FlashKDA provides the backend implementation in csrc/flash_kda.cpp, while FLA handles the dispatch decision and logging.

Frequently Asked Questions

How do I know if FlashKDA is being used instead of Triton?

Check your logs for the message [FLA Backend] kda.chunk_kda -> flashkda. If present, FlashKDA’s CUDA kernel in csrc/flash_kda.cpp is active. If you see rejected: instead, FLA has fallen back to its Triton implementation.

What does "unsupported CUDA arch" mean in the logs?

This indicates your GPU does not meet the SM90+ architecture requirement necessary for FlashKDA’s optimized kernels. You must run on compatible hardware (e.g., H100, A100) or use the Triton fallback by setting FLA_FLASH_KDA=0.

Can I enable logs without modifying my script?

Yes. You can pass the LOGLEVEL environment variable when running your script: LOGLEVEL=INFO python your_script.py. Alternatively, create a .py startup script that configures logging before your main code executes.

Where is the dispatch logic that generates these logs?

The dispatch logic resides in the flash-linear-attention (FLA) repository, not in FlashKDA itself. FlashKDA only provides the backend implementation via flash_kda/__init__.py and the CUDA kernels. FLA’s wrapper checks availability and emits the [FLA Backend] messages according to its internal selection criteria.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}