Add more tests

[drewjin]

Summary by CodeRabbit

• Tests

  • Added validation benchmark test suite with configuration files for comprehensive model behavior testing
  • Enhanced kernel testing with diagnostic visualization and debugging utilities
  • Implemented attention output validation against reference implementations
  • Removed legacy benchmark test module

• Chores

  • Updated ignore patterns to exclude build artifacts and output directories

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ae93ac5a-1e8f-4d8d-84a6-a1034716a39a

📥 Commits

Reviewing files that changed from the base of the PR and between d360800 and 86388e2.

📒 Files selected for processing (12)

• .gitignore
• diffulex_kernel/python/chunked_prefill_cutedsl.py
• diffulex_kernel/python/chunked_prefill_triton.py
• test/python/engine/config/validation_bench_sdar_bufsz1.yml
• test/python/engine/config/validation_bench_sdar_bufsz2.yml
• test/python/engine/config/validation_bench_sdar_bufsz4.yml
• test/python/engine/dummy_attn_with_validation.py
• test/python/engine/test_diffulex_bench.py
• test/python/engine/test_monkey_patch.py
• test/python/engine/test_validation_bench.py
• test/python/kernel/debug_chunked_prefill.py
• test/python/kernel/test_dllm_flash_attn_chunked_prefill_unified_kernel.py

---

📝 Walkthrough

Walkthrough

The PR introduces attention output validation infrastructure and benchmarking utilities. Changes include new YAML configuration files for validation benchmarks with varying buffer sizes, a validation wrapper for attention mechanisms, enhanced kernel testing with mask visualization capabilities, monkey-patching utilities for test hooks, and removal of legacy test code.

Changes

Cohort / File(s)	Summary
Configuration & Metadata .gitignore, test/python/engine/config/validation_bench_sdar_bufsz*.yml (×3)	Updated .gitignore to ignore cutlass and output/ directories. Added three YAML config files defining engine (model paths, decoding strategy multi_bd, KV cache layout, block/buffer sizing) and eval (GSM8K dataset, temperature, output settings) with varying buffer_size values (1, 2, 4).
Validation Wrapper test/python/engine/dummy_attn_with_validation.py	Introduces AttentionWithValidation class wrapping the original Attention with configurable tolerance and validation logic. Computes PyTorch reference outputs, builds block-causal masks, compares via torch.testing.assert_close, and logs detailed mismatches. Includes install_validation_hook() for monkey-patching.
Validation Test Infrastructure test/python/engine/test_monkey_patch.py, test/python/engine/test_validation_bench.py	Adds validation hook installation test and parametrized benchmark runner. test_validation_bench.py executes benchmarks in subprocesses with installed validation hooks across three config files, capturing and asserting subprocess exit codes.
Kernel Testing & Debugging test/python/kernel/test_dllm_flash_attn_chunked_prefill_unified_kernel.py, test/python/kernel/debug_chunked_prefill.py	Enhanced kernel tests with _visualize_mask() for attention-mask diagnostics and structured compact/full rendering. Changed default BLOCK_M from 128 to 64. Modified validation to iterate over all sequences. Added visualization test. New debug script for standalone kernel testing without pytest.
Documentation diffulex_kernel/python/chunked_prefill_triton.py	Added inline comment above @triton.autotune(...) decorator recommending disabling autotuner during testing to avoid slowdown.
Removed test/python/engine/test_diffulex_bench.py	Deleted legacy test module (240 lines) that configured pytest, resolved checkpoint paths, and ran parameterized Diffulex generation tests with multiple strategies.

Sequence Diagram(s)

sequenceDiagram
    participant Test as Test Code
    participant Wrapper as AttentionWithValidation
    participant Original as Original Attention
    participant PyTorch as PyTorch Reference
    participant Compare as Comparison

    Test->>Wrapper: forward(q, k, v, mask)
    activate Wrapper
    
    Wrapper->>Original: super().forward(...)
    activate Original
    Original-->>Wrapper: kernel_output
    deactivate Original
    
    Wrapper->>Wrapper: _compute_reference(q, k, v, metadata)
    activate Wrapper
    Wrapper->>PyTorch: scaled_dot_product_attention(...)
    activate PyTorch
    PyTorch-->>Wrapper: ref_output
    deactivate PyTorch
    deactivate Wrapper
    
    Wrapper->>Compare: torch.testing.assert_close
    activate Compare
    alt Output matches within tolerance
        Compare-->>Wrapper: ✓ Validation passed
    else Output exceeds tolerance
        Compare-->>Wrapper: ✗ Assertion error
        Wrapper->>Wrapper: Log error details
        Wrapper-->>Test: ✗ Raise exception
    end
    deactivate Compare
    
    deactivate Wrapper
    Wrapper-->>Test: result

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐰 A kernel wrapped in gentle care,
With validations everywhere,
Reference checks and masks so bright,
Attention outputs burning right!
Benchmarks run in subprocess dance,
While buffer sizes find their chance. ✨

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

Tip

CodeRabbit can use OpenGrep to find security vulnerabilities and bugs across 17+ programming languages.

OpenGrep is compatible with Semgrep configurations. Add an opengrep.yml or semgrep.yml configuration file to your project to enable OpenGrep analysis.