[Example] Add sparse mla examples

[LeiWang1999]

as title.

Summary by CodeRabbit

• New Features

  • Added FP8 MQA logits example with reference path, verification utilities, and performance benchmarking.
  • Introduced sparse multi-head attention forward examples (standard and pipelined) featuring compiled kernels, correctness checks against a reference, and throughput metrics.
  • Included a utilities module for sequence-length handling, lightweight caching, FP8 casting, and helper functions shared across examples.

• Documentation

  • Updated README to include a clear directory structure overview.

[coderabbitai]

Caution

Review failed

The pull request is closed.

Walkthrough

Adds a README directory structure. Introduces new FP8 MQA logits pipeline with JIT kernels, reference implementation, validation, and benchmarking. Adds two sparse MLA forward kernels (standard and pipelined) with interfaces, references, tests, and profiling. Provides a new utilities module for sequence-length math, FP8 casting, caching, and helpers.

Changes

Cohort / File(s)	Summary
Docs Update examples/deepseek_v32/README.md	Replaces placeholder with a directory tree describing files and roles.
FP8 MQA Logits Pipeline examples/deepseek_v32/fp8_mqa_logits.py	New FP8 MQA logits workflow: error/validation utilities, config generator, lazy tensor supplier, TileLang JIT kernels (fused logits and cleanup), public interface, FP32 reference, test/verification, and benchmarking.
Sparse MLA Forward Kernels examples/deepseek_v32/sparse_mla_fwd.py, examples/deepseek_v32/sparse_mla_fwd_pipelined.py	Two TileLang JIT sparse MLA forward implementations (standard and pipelined) with interfaces, reference PyTorch paths, input validation, masking/normalization, tests, and profiling.
Utilities for Seqlen/FP8/Caching examples/deepseek_v32/utils.py	New helpers: single-entry tensor cache, seq-index and cu_seqlen derivations, ks/ke computation, FP8 per-dim casting, random cu_seqlens generator, similarity checks, and CLI benchmark.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor User
  participant Py as Python Script
  participant TL as TileLang Kernel
  participant Ref as Reference Impl
  participant Utils as Utils

  User->>Py: Run fp8_mqa_logits.py (__main__)
  Py->>Utils: Generate cu_seqlens / scales
  Py->>TL: mqa_attn_return_logits_interface(q_fp8, kv_fp8, ...)
  TL-->>Py: logits_tl
  Py->>Ref: ref_fp8_mqa_logits(q_fp32, kv_fp32, ...)
  Ref-->>Py: logits_ref, cost_ref
  Py->>Py: validate_tensor_match(logits_tl, logits_ref)
  Py-->>User: Print diff, perf, cost

Loading

sequenceDiagram
  autonumber
  actor User
  participant Py as sparse_mla_fwd*_interface
  participant TL as TileLang Kernel
  participant Ref as PyTorch Ref
  participant Prof as Profiler

  User->>Py: Call interface(q, kv, indices, ...)
  Py->>Py: Validate shapes, compute sm_scale
  Py->>TL: Launch kernel (sparse_mla_fwd / pipelined)
  TL-->>Py: output, lse
  Py->>Ref: Optional reference run
  Ref-->>Py: output_ref
  Py->>Py: Compare outputs (optional)
  Py->>Prof: Benchmark
  Prof-->>User: Throughput/metrics

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

• tzj-fxz

Poem

I twitch my ears at kernels new,
FP8 sparks in midnight blue.
Sparse paths weave where queries tread,
Tiled lanes hum, softmax spread.
With cached carrots (tensors) near,
I benchmark, nibble, never fear—
Hop! The logits land so clear. 🥕🐇

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• 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 78664e2 and 4b804cf.

📒 Files selected for processing (5)

• examples/deepseek_v32/README.md (1 hunks)
• examples/deepseek_v32/fp8_mqa_logits.py (1 hunks)
• examples/deepseek_v32/sparse_mla_fwd.py (1 hunks)
• examples/deepseek_v32/sparse_mla_fwd_pipelined.py (1 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.}

[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!

🚀