[Example] Add correctness assert into dsa example

[LeiWang1999]

Summary by CodeRabbit

• New Features
  • Added built-in output validation using a global similarity metric with clear pass/fail feedback and color-coded messages, including graceful handling of all-zero cases.
• Refactor
  • Replaced element-wise comparisons with a global similarity score for faster, more robust result checks.
• Chores
  • Reduced the default example sequence size to 4096 to speed up runs and make quick local verification easier.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work!

🚀

[coderabbitai]

Caution

Review failed

The pull request is closed.

Walkthrough

Introduces a global tensor-similarity utility replacing prior helpers and integrates it into the sparse MLA forward example. The example now conditionally validates outputs against a reference when SKV <= 4096 and updates the default SKV to 4096.

Changes

Cohort / File(s)	Summary of Changes
Sparse MLA test harness examples/deepseek_v32/sparse_mla_fwd.py	Imported assert_tensors_similar; added conditional path to compute reference output and assert similarity when SKV ≤ 4096; prints confirmation on pass; changed default SKV from 32768 to 4096.
Utils: tensor similarity examples/deepseek_v32/utils.py	Replaced calc_sim/assert_similar with calculate_tensor_similarity/assert_tensors_similar using normalized dot-product similarity; added docstrings; color-coded logging for zero-magnitude, pass/fail; maintains special handling for all-zero tensors.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant U as User
  participant Main as sparse_mla_fwd.py (__main__)
  participant Kern as sparse_mla_fwd
  participant Ref as ref_sparse_mla_fwd_interface
  participant Util as assert_tensors_similar

  U->>Main: Run with SKV (default 4096)
  Main->>Kern: Compute tl_out
  alt SKV ≤ 4096
    Main->>Ref: Compute ref_out
    Main->>Util: Compare tl_out vs ref_out (eps)
    opt similarity < threshold
      Util-->>Main: Error (diff, threshold)
    end
    opt similarity ≥ threshold
      Util-->>Main: Pass confirmation
    end
  else SKV > 4096
    Main-->>U: Skip reference comparison
  end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

A rabbit taps keys with a whiskered grin,
Swaps old checks for one sleek cosine spin.
SKV shrinks—4096’s neat,
Reference agrees, results concrete.
Logs glow amber, red, or green—
Hop-hop, the tensors stay serene. 🐇✨

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 242cb45 and d94a226.

📒 Files selected for processing (2)

• examples/deepseek_v32/sparse_mla_fwd.py (3 hunks)
• examples/deepseek_v32/utils.py (1 hunks)

---

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.}