[BUG FIX] Improve error reporting and occupancy in benchmarks

[LoserCheems]

Summary

• This update enhances error reporting in benchmark tests and improves launch configuration logic for better occupancy in forward and sparse kernels.

Root Cause

• The previous error handling did not provide sufficient detail, making debugging difficult. Additionally, the occupancy settings for certain configurations were suboptimal.

Changes

• Added detailed error reporting with traceback in benchmark tests.
• Updated launch configuration logic to adjust tile sizes based on qheads_per_kvhead for improved occupancy.

Reproduction

• Run benchmark tests with various configurations to observe error messages and occupancy metrics.

Tests

• Validated changes by running existing benchmark tests and confirming improved error output and occupancy metrics.

Compatibility

• No backward compatibility issues identified.

Checklist

• Linked issue provided
• Adds or updates tests
• Updates docs if needed
• No perf regressions

[Copilot]

Pull request overview

This PR improves benchmark failure diagnostics by including full Python tracebacks in error messages, and updates Triton forward-kernel launch configuration heuristics intended to improve occupancy (especially for split-KV / decode-like cases).

Changes:

• Include traceback details in BenchmarkResult.error_message and print multi-line failure output in benchmark scripts.
• Adjust split-KV TILE_M selection based on qheads_per_kvhead and retune some SM90 (H100) forward launch configs.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
tests/benchmark_forward.py	Adds traceback capture to benchmark failures and prints full error details.
tests/benchmark_decode.py	Same traceback-enhanced error reporting for decode benchmarks.
tests/benchmark_backward.py	Same traceback-enhanced error reporting for backward benchmarks.
flash_sparse_attn/ops/triton/launch_template.py	Updates forward (dense/sparse/gated) launch config heuristics, including split-KV TILE_M and SM90 configs.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.