feat: Reland support lse in trtllm paged attn kernels (#3116)

<!-- .github/pull_request_template.md -->

## 📌 Description

Re-land https://github.com/flashinfer-ai/flashinfer/pull/3058, after
revert https://github.com/flashinfer-ai/flashinfer/pull/3079. Issues
with original land investigated in
https://github.com/flashinfer-ai/flashinfer/issues/3114.

Re-introduce compute + output LSE support to paged context / decode /
mla operators.

To validate against the original versions issues, I ran
`compute-sanitizer --tool memcheck` on the previously failing
prefill/decode/MLA cases reports with 0 errors (prefill
`[...-256-32-2-8-HND]`), confirming the OOB is fixed.

Changes from original version:

**C++ changes**:
`include/flashinfer/trtllm/fmha/fmhaRunnerParams.h`: Added
`lseStrideTokens` and `lseStrideHeads` next to `lsePtr`.
`csrc/trtllm_fmha_kernel_launcher.cu`:
`trtllm_paged_attention_launcher` / `trtllm_ragged_attention_launcher`:
made `softmaxStatsPtr` allocation conditional on `lse != nullptr`, sized
it to the kernel's tile-aware layout `num_qo_heads * batch_size *
round_up(max_q_len, 256) * sizeof(float2)`, and plumbed LSE pointer +
strides into `runner_params`.
Kept the 8MB counter slab at a fixed offset in generation mode so the
existing workspace guard assertion still holds.
`trtllm_paged_attention_decode` / `trtllm_paged_attention_context` /
`trtllm_ragged_attention` wrappers now accept `Optional<TensorView> lse`
plus strides and forward them to the launcher.
**Python API**
`flashinfer/decode.py` (`trtllm_batch_decode_with_kv_cache`,
`TrtllmGenDecodeModule._paged_run`), `flashinfer/mla/_core.py`
(`trtllm_batch_decode_with_kv_cache_mla`), `flashinfer/prefill.py`
(`trtllm_batch_context_with_kv_cache`, `_paged_run in
get_trtllm_gen_prefill_module`): added `lse` / `return_lse` args with
shape `(num_tokens, num_qo_heads)` `float32` validation, passed strides
through.
XQA (and cute-dsl for MLA) backends raise `NotImplementedError` when
`return_lse=True` / `lse is not None`.

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Optional LSE (Log-Sum-Exp) buffer support across paged/ragged
attention and public decode/context APIs; callers can provide or request
returned LSE (trtllm-gen only). Stride metadata is computed and
propagated; unsupported backends raise NotImplemented when LSE is
requested.

* **Bug Fixes / Stability**
* Runtime LSE validation (dtype/shape), conditional workspace
allocation/initialization to avoid unnecessary buffers, and guard-byte
protection to detect/prevent out-of-bounds writes.

* **Tests**
* End-to-end LSE tests validating dtype/shape/values and guard-region
overflow detection.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Alex Yang <aleozlx@gmail.com>