Logging¶

# This works correctly - no synchronization errors
with torch.cuda.graph(cuda_graph):
    result = mm_fp4(a, b, scales, ...)  # Level 5 logging active
    # Statistics automatically skipped during capture

"""
Producer script: Run bmm_fp8 with Level 10 logging to dump tensors.

Usage:
    FLASHINFER_LOGLEVEL=10 FLASHINFER_DUMP_DIR=./bmm_fp8_dumps python bmm_fp8_producer.py
"""
import torch
from flashinfer import bmm_fp8

def to_float8(x, dtype=torch.float8_e4m3fn):
    """Convert tensor to FP8 with per-tensor scaling."""
    finfo = torch.finfo(dtype)
    min_val, max_val = x.aminmax()
    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
    scale = finfo.max / amax
    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
    return x_scl_sat.to(dtype), scale.float().reciprocal()

# Parameters
b, m, n, k = 4, 64, 128, 256
input_dtype = torch.float8_e4m3fn
mat2_dtype = torch.float8_e4m3fn
res_dtype = torch.bfloat16

# Create input tensors
input_bf16 = torch.randn([b, m, k], device="cuda", dtype=torch.bfloat16)
input_fp8, input_inv_s = to_float8(input_bf16, dtype=input_dtype)

# mat2: row major -> column major (transposed)
mat2_bf16 = torch.randn([b, n, k], device="cuda", dtype=torch.bfloat16).transpose(-2, -1)
mat2_fp8, mat2_inv_s = to_float8(mat2_bf16, dtype=mat2_dtype)

# Pre-allocate output
res = torch.empty([b, m, n], device="cuda", dtype=res_dtype)

# Call bmm_fp8 - this will be logged/dumped at Level 10
bmm_fp8(input_fp8, mat2_fp8, input_inv_s, mat2_inv_s, res_dtype, res, backend="cublas")

# Print a small portion of the output for verification
print("Output shape:", res.shape)
print("Output[0, :3, :3]:")
print(res[0, :3, :3])

"""
Reproducer script: Load dumped tensors and replay bmm_fp8.

Usage:
    python bmm_fp8_reproducer.py
"""
import torch
from pathlib import Path
from flashinfer import bmm_fp8
from flashinfer.api_logging import replay_from_dump

DUMP_DIR = "./bmm_fp8_dumps"

# Find the bmm_fp8 dump directory (should be the only one or the latest)
dump_path = Path(DUMP_DIR)
bmm_dumps = sorted([d for d in dump_path.iterdir() if d.is_dir() and "bmm_fp8" in d.name])
latest_dump = bmm_dumps[-1]  # Use the latest dump
print(f"Loading dump from: {latest_dump}")

# Use replay_from_dump to load inputs and optionally execute
result = replay_from_dump(
    str(latest_dump),
    compare_outputs=True,  # Load expected outputs for comparison
    device="cuda",
    run=False,  # We'll call the function manually below
)

# Extract the loaded arguments - args contains all positional args including the output tensor
args = result["args"]
kwargs = result["kwargs"]
expected_tensors = result.get("expected_tensors", {})

# Replay the call - args already contains (input, mat2, input_inv_s, mat2_inv_s, dtype, out)
res = bmm_fp8(*args, **kwargs)

# Print the same portion for comparison
print("Replayed output shape:", res.shape)
print("Replayed output[0, :3, :3]:")
print(res[0, :3, :3])

# Compare with expected output if available
if "result" in expected_tensors:
    expected = expected_tensors["result"]
    if torch.allclose(res, expected, rtol=1e-3, atol=1e-3):
        print("\n✅ Output matches expected result!")
    else:
        diff = (res - expected).abs().max().item()
        print(f"\n❌ Output mismatch! Max diff: {diff}")

"""
Producer script: Run BatchDecodeWithPagedKVCacheWrapper with Level 10 logging.

Usage:
    FLASHINFER_LOGLEVEL=10 FLASHINFER_DUMP_DIR=./batch_decode_dumps python batch_decode_producer.py
"""
import torch
import flashinfer

# Parameters
batch_size = 4
kv_len = 512
page_size = 16
num_kv_heads = 4
num_qo_heads = 32
head_dim = 128
kv_layout = "NHD"

# Create query tensor
q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda", dtype=torch.float16)

# Create paged KV cache
num_pages_per_seq = (kv_len + page_size - 1) // page_size
total_num_pages = num_pages_per_seq * batch_size
kv_shape = [total_num_pages, 2, page_size, num_kv_heads, head_dim]  # NHD layout
kv_data = torch.randn(*kv_shape, device="cuda", dtype=torch.float16)

# Create index tensors
kv_indptr = torch.arange(0, batch_size + 1, device="cuda", dtype=torch.int32) * num_pages_per_seq
kv_indices = torch.arange(0, total_num_pages, device="cuda", dtype=torch.int32)
kv_last_page_len = torch.full(
    (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda"
)

# Create workspace and wrapper - __init__ will be logged
workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.int8, device="cuda")
wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(workspace_buffer, kv_layout)

# Plan - will be logged
wrapper.plan(
    kv_indptr,
    kv_indices,
    kv_last_page_len,
    num_qo_heads,
    num_kv_heads,
    head_dim,
    page_size,
    data_type=torch.float16,
    q_data_type=torch.float16,
)

# Run - will be logged
output, lse = wrapper.run(q, kv_data, return_lse=True)

# Print a small portion of the output
print("Output shape:", output.shape)
print("Output[0, :3, :3]:")
print(output[0, :3, :3])
print("\nLSE shape:", lse.shape)
print("LSE[0, :5]:", lse[0, :5])

"""
Reproducer script: Replay BatchDecodeWithPagedKVCacheWrapper calls.

Usage:
    python batch_decode_reproducer.py
"""
import torch
from pathlib import Path
from flashinfer.api_logging import replay_sequence

DUMP_DIR = "./batch_decode_dumps"

# replay_sequence handles stateful objects automatically via object_registry
# It will:
# 1. Replay __init__ to create the wrapper instance
# 2. Replay plan() on the same instance
# 3. Replay run() on the same instance and compare outputs
results = replay_sequence(DUMP_DIR, device="cuda")

# Print summary
passed = 0
failed = 0
for i, res in enumerate(results):
    func_name = res.get("metadata", {}).get("function_name", "unknown")
    dump_dir = Path(res.get("dump_dir", "")).name

    if "error" in res:
        print(f"[{i+1}] {func_name} ({dump_dir}): ❌ Error: {res['error']}")
        failed += 1
    elif res.get("comparison_match", True):
        print(f"[{i+1}] {func_name} ({dump_dir}): ✅ Passed")
        passed += 1
    else:
        print(f"[{i+1}] {func_name} ({dump_dir}): ❌ Mismatch")
        failed += 1

print(f"\nSummary: {passed} passed, {failed} failed")

# For manual inspection, you can also access individual results
# Find the 'run' call result (usually the last non-init, non-plan call)
for res in results:
    func_name = res.get("metadata", {}).get("function_name", "")
    if "run" in func_name and "execution_result" in res:
        output = res["execution_result"]
        if isinstance(output, tuple):
            output_tensor, lse = output
            print("\nReplayed output[0, :3, :3]:")
            print(output_tensor[0, :3, :3])
            print("Replayed LSE[0, :5]:", lse[0, :5])
        break

"""
Manual replay: Load tensors directly from .pt files.
"""
import json
import torch
from pathlib import Path
from flashinfer import bmm_fp8

# Path is an example, replace with the actual path.
dump_dir = Path("./bmm_fp8_dumps/20250108_103217_012_pid12345_bmm_fp8_call0001")

# Load metadata from JSONL (read last line for most complete state)
with open(dump_dir / "metadata.jsonl") as f:
    lines = [line.strip() for line in f if line.strip()]
    metadata = json.loads(lines[-1])  # Last line has completed state

print(f"Function: {metadata['function_name']}")
print(f"Module: {metadata['module']}")
print(f"Status: {metadata['execution_status']}")
print(f"Input tensors: {metadata['tensor_info']['input_tensor_keys']}")

# Load input tensors
inputs = torch.load(dump_dir / "inputs.pt", map_location="cuda")

# Load expected outputs (if execution completed successfully)
outputs_path = dump_dir / "outputs.pt"
if outputs_path.exists():
    expected = torch.load(outputs_path, map_location="cuda")
    print(f"Output tensors: {list(expected.keys())}")

# Tensors are ready to use - reconstruct the call as needed
for key, tensor in inputs.items():
    print(f"  {key}: shape={tensor.shape}, dtype={tensor.dtype}")

"""
Scan session.jsonl for quick overview of recorded calls.
"""
import json
from pathlib import Path
from collections import Counter

dump_root = Path("./my_dumps")
session_file = dump_root / "session.jsonl"

# Read all records
records = []
with open(session_file) as f:
    for line in f:
        if line.strip():
            records.append(json.loads(line))

# Filter to completed calls only
completed = [r for r in records if r["execution_status"] == "completed"]
print(f"Total completed calls: {len(completed)}")

# Count by function name
func_counts = Counter(r["function_name"] for r in completed)
print("\nCalls by function:")
for func, count in func_counts.most_common():
    print(f"  {func}: {count}")

# Find calls that didn't complete (potential crashes)
inputs_only = [r for r in records if r["execution_status"] == "inputs_saved"]
# Group by dump_dir to find incomplete calls
completed_dirs = {r["dump_dir"] for r in completed}
incomplete = [r for r in inputs_only if r["dump_dir"] not in completed_dirs]
if incomplete:
    print(f"\n⚠️  Found {len(incomplete)} incomplete calls (potential crashes):")
    for r in incomplete:
        print(f"  - {r['function_name']} at {r['dump_dir']}")