reference/code/Diffulex-mbd-lms/test/python/utils/checker.py

def CHECK_D2F_SLOT_MAPPING(seqs, slot_mapping):
    # check slot mapping layout
    start_idx = 0
    for seq in seqs:
        cur_ref_slot_mapping = []
        for idx in range(seq.num_diffusion_blocks):
            if seq.active_blocks[idx]:
                padding_num_tokens = (seq.num_diffusion_blocks - idx) * seq.diffusion_block_size
                cur_ref_slot_mapping.extend([-1] * padding_num_tokens)
                break
            elif seq.to_cache_blocks[idx]:
                cur_ref_slot_mapping.extend([0] * seq.diffusion_block_size)
        cur_slot_mapping = slot_mapping[start_idx : start_idx + len(cur_ref_slot_mapping)]
        for slot, ref_slot in zip(cur_slot_mapping, cur_ref_slot_mapping):
            try:
                if ref_slot == -1:
                    assert slot == -1
                elif ref_slot == 0:
                    assert slot != -1
                elif ref_slot is not None:
                    assert slot is not None
            except AssertionError:
                raise ValueError(
                    f"Slot mapping mismatch: {slot} != {ref_slot}. "
                    f"Check the implementation of prepare_decode.\n"
                    f"slot_mapping: {cur_slot_mapping}\n"
                    f"ref_slot_mapping: {cur_ref_slot_mapping}\n"
                    f"diff: {[s - r for s, r in zip(cur_slot_mapping, cur_ref_slot_mapping)]}"
                )
        start_idx += len(cur_ref_slot_mapping)


def CHECK_FLASH_ATTN_PREFILL(
    q,
    k,
    v,
    cu_seqlens_q,
    cu_seqlens_k,
    max_seqlen_q,
    prefill_kernel,
    diffusion_block_size: int = 32,
    is_block_attn: bool = False,
):
    """
    Verify prefill kernel correctness by comparing with PyTorch's scaled_dot_product_attention.

    Args:
        q: Query tensor [total_q_len, num_heads, head_dim]
        k: Key tensor [total_kv_len, num_kv_heads, head_dim]
        v: Value tensor [total_kv_len, num_kv_heads, head_dim]
        cu_seqlens_q: Cumulative sequence lengths for queries
        cu_seqlens_k: Cumulative sequence lengths for keys/values
        max_seqlen_q: Maximum sequence length for queries
        prefill_kernel: The kernel function to test
        diffusion_block_size: Size of diffusion blocks for block attention
        is_block_attn: Whether this is block attention mode
    """
    import torch
    import torch.nn.functional as F
    from einops import rearrange

    # Run kernel
    kernel_output = prefill_kernel(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q)

    # Compute reference output using PyTorch's SDPA
    head_dim = q.shape[2]
    scale = 1.0 / (head_dim**0.5)
    num_seqs = len(cu_seqlens_q) - 1

    gt_output = torch.zeros_like(q)
    for seq_idx in range(num_seqs):
        q_start = cu_seqlens_q[seq_idx].item()
        q_end = cu_seqlens_q[seq_idx + 1].item()
        kv_start = cu_seqlens_k[seq_idx].item()
        kv_end = cu_seqlens_k[seq_idx + 1].item()

        q_seq = q[q_start:q_end]
        k_seq = k[kv_start:kv_end]
        v_seq = v[kv_start:kv_end]

        q_len = q_seq.shape[0]
        kv_len = k_seq.shape[0]

        # Reshape for SDPA: [1, num_heads, seq_len, head_dim]
        q_sdpa = rearrange(q_seq, "s h d -> 1 h s d")
        k_sdpa = rearrange(k_seq, "s h d -> 1 h s d")
        v_sdpa = rearrange(v_seq, "s h d -> 1 h s d")

        if not is_block_attn:
            # Standard attention
            attn_out = F.scaled_dot_product_attention(
                q_sdpa,
                k_sdpa,
                v_sdpa,
                dropout_p=0.0,
                is_causal=False,
                scale=scale,
                enable_gqa=True,
            )
        else:
            # Block attention with mask
            block_mask = torch.zeros((1, 1, q_len, kv_len), dtype=q.dtype, device=q.device).bool()
            num_diffusion_blocks = (kv_len + diffusion_block_size - 1) // diffusion_block_size
            for block_idx in range(num_diffusion_blocks):
                block_start = block_idx * diffusion_block_size
                block_end = min(block_start + diffusion_block_size, kv_len)
                block_mask[..., block_start:block_end, :block_end] = True

            attn_out = F.scaled_dot_product_attention(
                q_sdpa,
                k_sdpa,
                v_sdpa,
                attn_mask=block_mask,
                dropout_p=0.0,
                is_causal=False,
                scale=scale,
                enable_gqa=True,
            )

        gt_output[q_start:q_end] = rearrange(attn_out, "1 h s d -> s h d").to(gt_output.dtype)

    # Compare results
    atol = 1e-2
    rtol = 1e-2
    try:
        torch.testing.assert_close(
            kernel_output,
            gt_output,
            atol=atol,
            rtol=rtol,
            msg="Kernel output does not match reference implementation",
        )
    except AssertionError as e:
        # Compute error statistics for debugging
        abs_diff = torch.abs(kernel_output - gt_output)
        max_diff = torch.max(abs_diff).item()
        mean_diff = torch.mean(abs_diff).item()
        rel_diff = torch.abs((kernel_output - gt_output) / (gt_output + 1e-8))
        max_rel_diff = torch.max(rel_diff).item()
        mean_rel_diff = torch.mean(rel_diff).item()

        # Count elements that exceed tolerance
        total_elements = kernel_output.numel()
        # Elements that exceed absolute tolerance
        exceeds_atol = abs_diff > atol
        num_exceeds_atol = exceeds_atol.sum().item()
        # Elements that exceed relative tolerance
        exceeds_rtol = rel_diff > rtol
        num_exceeds_rtol = exceeds_rtol.sum().item()
        # Elements that exceed either tolerance
        exceeds_tolerance = exceeds_atol | exceeds_rtol
        num_exceeds_tolerance = exceeds_tolerance.sum().item()
        pct_exceeds_tolerance = (num_exceeds_tolerance / total_elements * 100) if total_elements > 0 else 0

        raise AssertionError(
            f"Prefill kernel verification failed!\n"
            f"Max absolute difference: {max_diff:.6f}\n"
            f"Mean absolute difference: {mean_diff:.6f}\n"
            f"Max relative difference: {max_rel_diff:.6f}\n"
            f"Mean relative difference: {mean_rel_diff:.6f}\n"
            f"Total elements: {total_elements}\n"
            f"Elements exceeding absolute tolerance (atol={atol}): {num_exceeds_atol} ({num_exceeds_atol / total_elements * 100:.2f}%)\n"
            f"Elements exceeding relative tolerance (rtol={rtol}): {num_exceeds_rtol} ({num_exceeds_rtol / total_elements * 100:.2f}%)\n"
            f"Elements exceeding either tolerance: {num_exceeds_tolerance} ({pct_exceeds_tolerance:.2f}%)\n"
            f"Kernel output shape: {kernel_output.shape}\n"
            f"Reference output shape: {gt_output.shape}\n"
            f"Original error: {str(e)}"
        )


def CHECK_FLASH_ATTN_DECODE(
    q,
    k,
    v,
    k_cache,
    v_cache,
    block_tables,
    context_lens,
    cu_seqlens_q,
    cu_seqlens_k,
    max_seqlen_q,
    decode_kernel,
    scale: float,
    num_groups: int,
    page_block_size: int,
    diffusion_block_size: int = 32,
    is_block_attn: bool = False,
):
    """
    Verify decode kernel correctness by comparing with reference implementation.
    This function mimics engine-like scenarios with memory reuse testing.

    Args:
        q: Query tensor [total_q_len, num_heads, head_dim]
        k: Key tensor [total_kv_len, num_kv_heads, head_dim]
        v: Value tensor [total_kv_len, num_kv_heads, head_dim]
        k_cache: KV cache for keys [num_page_blocks, page_block_size, num_kv_heads, head_dim]
        v_cache: KV cache for values [num_page_blocks, page_block_size, num_kv_heads, head_dim]
        block_tables: Block tables [num_seqs, max_seq_num_blocks]
        context_lens: Context lengths for each sequence [num_seqs]
        cu_seqlens_q: Cumulative sequence lengths for queries
        cu_seqlens_k: Cumulative sequence lengths for keys/values
        max_seqlen_q: Maximum sequence length for queries
        decode_kernel: The kernel function to test
        scale: Attention scale factor
        num_groups: Number of GQA groups (num_heads // num_kv_heads)
        page_block_size: Size of page blocks in KV cache
        diffusion_block_size: Size of diffusion blocks for block attention
        is_block_attn: Whether this is block attention mode
    """
    import torch
    from test.python.kernel.test_dllm_flash_attn_decode_kernel import (
        naive_sdpa_with_kv_cache,
    )

    # Test with memory reuse (simulate engine's behavior)
    # Run multiple times to check for memory corruption
    outputs = []
    for run_idx in range(3):
        output = decode_kernel(
            q,
            k,
            v,
            k_cache,
            v_cache,
            block_tables,
            context_lens,
            cu_seqlens_q,
            cu_seqlens_k,
            max_seqlen_q,
        )
        torch.cuda.synchronize()
        outputs.append(output.clone())

    # Verify consistency across runs
    consistent = True
    for i in range(1, len(outputs)):
        if not torch.allclose(outputs[0], outputs[i], atol=1e-5, rtol=1e-5):
            consistent = False
            max_diff = (outputs[0] - outputs[i]).abs().max().item()
            raise AssertionError(
                f"Output inconsistency detected in run {i}: max_diff={max_diff:.6f}. "
                f"This indicates potential memory corruption or non-deterministic behavior."
            )

    # Use the first output for comparison
    kernel_output = outputs[0]

    # Compute reference output using naive_sdpa_with_kv_cache (same as test file)
    gt_output = naive_sdpa_with_kv_cache(
        q,
        k,
        v,
        k_cache,
        v_cache,
        block_tables,
        context_lens,
        cu_seqlens_q,
        cu_seqlens_k,
        scale,
        num_groups,
        page_block_size,
    )

    # Compare results (using same tolerance as test file)
    atol = 1e-2
    rtol = 1e-2
    try:
        torch.testing.assert_close(
            kernel_output,
            gt_output,
            atol=atol,
            rtol=rtol,
            msg="Decode kernel output does not match reference implementation",
        )
    except AssertionError as e:
        # Compute error statistics for debugging
        abs_diff = torch.abs(kernel_output - gt_output)
        max_diff = torch.max(abs_diff).item()
        mean_diff = torch.mean(abs_diff).item()
        rel_diff = torch.abs((kernel_output - gt_output) / (gt_output + 1e-8))
        max_rel_diff = torch.max(rel_diff).item()
        mean_rel_diff = torch.mean(rel_diff).item()

        # Count elements that exceed tolerance
        total_elements = kernel_output.numel()
        # Elements that exceed absolute tolerance
        exceeds_atol = abs_diff > atol
        num_exceeds_atol = exceeds_atol.sum().item()
        # Elements that exceed relative tolerance
        exceeds_rtol = rel_diff > rtol
        num_exceeds_rtol = exceeds_rtol.sum().item()
        # Elements that exceed either tolerance
        exceeds_tolerance = exceeds_atol | exceeds_rtol
        num_exceeds_tolerance = exceeds_tolerance.sum().item()
        pct_exceeds_tolerance = (num_exceeds_tolerance / total_elements * 100) if total_elements > 0 else 0

        # Save test case data for debugging
        import os
        from pathlib import Path
        import pickle
        from datetime import datetime

        save_dir = Path(os.getenv("TEST_CASE_SAVE_DIR", "./failed_test_cases"))
        save_dir.mkdir(parents=True, exist_ok=True)

        # Generate unique filename with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        case_name = f"decode_kernel_failure_{timestamp}"
        case_dir = save_dir / case_name
        case_dir.mkdir(parents=True, exist_ok=True)

        # Save all input and output tensors
        test_data = {
            "inputs": {
                "q": q.cpu(),
                "k": k.cpu(),
                "v": v.cpu(),
                "k_cache": k_cache.cpu(),
                "v_cache": v_cache.cpu(),
                "block_tables": block_tables.cpu(),
                "context_lens": context_lens.cpu(),
                "cu_seqlens_q": cu_seqlens_q.cpu(),
                "cu_seqlens_k": cu_seqlens_k.cpu(),
            },
            "outputs": {
                "kernel_output": kernel_output.cpu(),
                "gt_output": gt_output.cpu(),
                "abs_diff": abs_diff.cpu(),
                "rel_diff": rel_diff.cpu(),
            },
            "parameters": {
                "max_seqlen_q": max_seqlen_q,
                "scale": scale,
                "num_groups": num_groups,
                "page_block_size": page_block_size,
                "diffusion_block_size": diffusion_block_size,
                "is_block_attn": is_block_attn,
                "atol": atol,
                "rtol": rtol,
            },
            "statistics": {
                "max_diff": max_diff,
                "mean_diff": mean_diff,
                "max_rel_diff": max_rel_diff,
                "mean_rel_diff": mean_rel_diff,
                "total_elements": total_elements,
                "num_exceeds_atol": num_exceeds_atol,
                "num_exceeds_rtol": num_exceeds_rtol,
                "num_exceeds_tolerance": num_exceeds_tolerance,
                "pct_exceeds_tolerance": pct_exceeds_tolerance,
            },
            "shapes": {
                "q_shape": list(q.shape),
                "k_shape": list(k.shape),
                "v_shape": list(v.shape),
                "k_cache_shape": list(k_cache.shape),
                "v_cache_shape": list(v_cache.shape),
                "block_tables_shape": list(block_tables.shape),
                "kernel_output_shape": list(kernel_output.shape),
                "gt_output_shape": list(gt_output.shape),
            },
        }

        # Save as pickle
        with open(case_dir / "test_data.pkl", "wb") as f:
            pickle.dump(test_data, f)

        # Save kernel source (same as test file)
        kernel_path = None
        try:
            kernel_source = decode_kernel.get_kernel_source()
            kernel_path = case_dir / "kernel.cu"
            kernel_path.write_text(kernel_source)
        except Exception:
            # If kernel source is not available, log but don't fail
            pass

        # Generate a Python script to reproduce the test case
        timestamp_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        repro_script = f'''"""
Auto-generated test case from failed CHECK_FLASH_ATTN_DECODE.
Generated at: {timestamp_str}

To use this test case:
1. Load the data: test_data = pickle.load(open("test_data.pkl", "rb"))
2. Move tensors to device: q = test_data['inputs']['q'].to(device), etc.
3. Call your kernel with the loaded inputs
"""
import torch
import pickle
from pathlib import Path

# Load test data
case_dir = Path(__file__).parent
with open(case_dir / "test_data.pkl", "rb") as f:
    test_data = pickle.load(f)

# Extract inputs
q = test_data['inputs']['q']
k = test_data['inputs']['k']
v = test_data['inputs']['v']
k_cache = test_data['inputs']['k_cache']
v_cache = test_data['inputs']['v_cache']
block_tables = test_data['inputs']['block_tables']
context_lens = test_data['inputs']['context_lens']
cu_seqlens_q = test_data['inputs']['cu_seqlens_q']
cu_seqlens_k = test_data['inputs']['cu_seqlens_k']

# Extract parameters
params = test_data['parameters']
max_seqlen_q = params['max_seqlen_q']
scale = params['scale']
num_groups = params['num_groups']
page_block_size = params['page_block_size']
diffusion_block_size = params['diffusion_block_size']
is_block_attn = params['is_block_attn']

# Extract expected outputs for comparison
gt_output = test_data['outputs']['gt_output']

# Print test case info
print("Test Case Information:")
q_shape = test_data['shapes']['q_shape']
k_shape = test_data['shapes']['k_shape']
v_shape = test_data['shapes']['v_shape']
print(f"  Shapes: q={{q_shape}}, k={{k_shape}}, v={{v_shape}}")
print(f"  Parameters: scale={{scale}}, num_groups={{num_groups}}, page_block_size={{page_block_size}}")
max_diff_val = test_data['statistics']['max_diff']
num_mismatches = test_data['statistics']['num_exceeds_tolerance']
print(f"  Statistics: max_diff={{max_diff_val:.6f}}, num_mismatches={{num_mismatches}}")

# TODO: Add your kernel call here
# kernel_output = your_kernel(q, k, v, k_cache, v_cache, block_tables, context_lens,
#                             cu_seqlens_q, cu_seqlens_k, max_seqlen_q)
# torch.testing.assert_close(kernel_output, gt_output, atol=params['atol'], rtol=params['rtol'])
'''

        with open(case_dir / "reproduce_test.py", "w") as f:
            f.write(repro_script)

        # Save error summary
        error_summary = f"""Test Case Failure Summary
Generated at: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

Shapes:
  q: {test_data["shapes"]["q_shape"]}
  k: {test_data["shapes"]["k_shape"]}
  v: {test_data["shapes"]["v_shape"]}
  k_cache: {test_data["shapes"]["k_cache_shape"]}
  v_cache: {test_data["shapes"]["v_cache_shape"]}
  block_tables: {test_data["shapes"]["block_tables_shape"]}
  kernel_output: {test_data["shapes"]["kernel_output_shape"]}
  gt_output: {test_data["shapes"]["gt_output_shape"]}

Parameters:
  max_seqlen_q: {max_seqlen_q}
  scale: {scale}
  num_groups: {num_groups}
  page_block_size: {page_block_size}
  diffusion_block_size: {diffusion_block_size}
  is_block_attn: {is_block_attn}
  atol: {atol}
  rtol: {rtol}

Statistics:
  Max absolute difference: {max_diff:.6f}
  Mean absolute difference: {mean_diff:.6f}
  Max relative difference: {max_rel_diff:.6f}
  Mean relative difference: {mean_rel_diff:.6f}
  Total elements: {total_elements}
  Elements exceeding absolute tolerance: {num_exceeds_atol} ({num_exceeds_atol / total_elements * 100:.2f}%)
  Elements exceeding relative tolerance: {num_exceeds_rtol} ({num_exceeds_rtol / total_elements * 100:.2f}%)
  Elements exceeding either tolerance: {num_exceeds_tolerance} ({pct_exceeds_tolerance:.2f}%)
"""

        with open(case_dir / "error_summary.txt", "w") as f:
            f.write(error_summary)

        save_info = f"\n\nTest case data saved to: {case_dir}\n"
        save_info += "  - test_data.pkl: All input/output tensors and metadata\n"
        save_info += "  - reproduce_test.py: Script to reproduce the test case\n"
        save_info += "  - error_summary.txt: Summary of the failure\n"
        if kernel_path is not None:
            save_info += "  - kernel.cu: CUDA kernel source code\n"

        # Show mismatched elements layout
        mismatch_info = ""
        if num_exceeds_tolerance > 0:
            # Get indices of mismatched elements
            mismatch_indices = torch.nonzero(exceeds_tolerance, as_tuple=False)
            num_to_show = min(50, num_exceeds_tolerance)  # Show at most 50 mismatches

            mismatch_info = f"\n\nMismatched elements (showing first {num_to_show} of {num_exceeds_tolerance}):\n"
            mismatch_info += "-" * 100 + "\n"
            mismatch_info += f"{'Index':<30} {'Kernel Value':<20} {'Ref Value':<20} {'Abs Diff':<15} {'Rel Diff':<15}\n"
            mismatch_info += "-" * 100 + "\n"

            for i in range(num_to_show):
                idx = mismatch_indices[i]
                idx_tuple = tuple(idx.tolist())

                kernel_val = kernel_output[idx_tuple].item()
                gt_val = gt_output[idx_tuple].item()
                abs_err = abs_diff[idx_tuple].item()
                rel_err = rel_diff[idx_tuple].item()

                mismatch_info += (
                    f"{str(idx_tuple):<30} {kernel_val:>19.6f} {gt_val:>19.6f} {abs_err:>14.6f} {rel_err:>14.6f}\n"
                )

            if num_exceeds_tolerance > num_to_show:
                mismatch_info += f"\n... and {num_exceeds_tolerance - num_to_show} more mismatches\n"

            # Show distribution of mismatches by dimension
            if len(kernel_output.shape) >= 2:
                mismatch_info += "\nMismatch distribution by dimensions:\n"
                for dim_idx in range(len(kernel_output.shape)):
                    dim_mismatches = exceeds_tolerance.sum(
                        dim=tuple(j for j in range(len(kernel_output.shape)) if j != dim_idx)
                    )
                    mismatch_info += (
                        f"  Dim {dim_idx} (size {kernel_output.shape[dim_idx]}): {dim_mismatches.tolist()}\n"
                    )

        raise AssertionError(
            f"Decode kernel verification failed!\n"
            f"Max absolute difference: {max_diff:.6f}\n"
            f"Mean absolute difference: {mean_diff:.6f}\n"
            f"Max relative difference: {max_rel_diff:.6f}\n"
            f"Mean relative difference: {mean_rel_diff:.6f}\n"
            f"Total elements: {total_elements}\n"
            f"Elements exceeding absolute tolerance (atol={atol}): {num_exceeds_atol} ({num_exceeds_atol / total_elements * 100:.2f}%)\n"
            f"Elements exceeding relative tolerance (rtol={rtol}): {num_exceeds_rtol} ({num_exceeds_rtol / total_elements * 100:.2f}%)\n"
            f"Elements exceeding either tolerance: {num_exceeds_tolerance} ({pct_exceeds_tolerance:.2f}%)\n"
            f"Kernel output shape: {kernel_output.shape}\n"
            f"Reference output shape: {gt_output.shape}\n"
            f"{mismatch_info}"
            f"{save_info}"
            f"Original error: {str(e)}"
        )