[MIRROR][Feature Request] Return LSE from all TRT LLM attention kernels

[yzh119]

Mirror of: flashinfer-ai#2169

Currently only some APIs return LSEs, while underlying kernels are capable of returning LSEs in all cases. Extend API to always return LSEs if asked. Example implementation: fw-ai/flashinfer#7

[yzh119]

@claude

[claude]

Claude finished @yzh119's task —— View job • claude/issue-7-20251207-2011 • Create PR ➔

---

✅ Completed Issue #7

Summary

Successfully implemented LSE (Log-Sum-Exp) return capability for all TRT LLM attention kernels. All underlying kernels were already capable of returning LSEs - this change extends the Python API to expose this functionality consistently across all functions.

Changes Made

Modified 5 functions across 3 files to support optional LSE returns:

1. trtllm_batch_context_with_kv_cache (flashinfer/prefill.py:3408)

• Context/prefill phase attention with paged KV cache
• Backend: trtllm-gen (Blackwell SM100/103)
• LSE shape: [num_tokens, num_heads]

2. trtllm_batch_decode_with_kv_cache (flashinfer/decode.py:2059)

• Decode phase attention with paged KV cache
• Backends: trtllm-gen (Blackwell) and xqa (Hopper/Blackwell)
• LSE shape: [num_tokens, num_heads]

3. xqa_batch_decode_with_kv_cache (flashinfer/decode.py:2374)

• Helper function for xqa backend
• LSE shape: [num_tokens, num_heads]

4. trtllm_batch_decode_with_kv_cache_mla (flashinfer/mla.py:514)

• Multi-head Latent Attention (MLA) for DeepSeek v3
• Backends: trtllm-gen (SM100/103) and xqa (SM120)
• LSE shape: [batch_size, num_heads]

5. xqa_batch_decode_with_kv_cache_mla (flashinfer/mla.py:720)

• Helper function for xqa MLA backend
• LSE shape: [batch_size, num_heads]

API Design

Each function now accepts two new optional parameters:

return_lse: bool = False          # Whether to return LSE
lse: Optional[torch.Tensor] = None  # Pre-allocated LSE tensor (optional)

Return behavior:

• When return_lse=False (default): Returns output tensor only (backward compatible)
• When return_lse=True: Returns tuple (output, lse)

LSE tensor:

• Dtype: torch.float32
• Auto-allocated if not provided
• Can be pre-allocated for memory efficiency

Implementation Pattern

Followed the existing pattern from trtllm_ragged_attention_deepseek to ensure API consistency.

Files Changed

• flashinfer/prefill.py - 52 lines added
• flashinfer/decode.py - 71 lines added
• flashinfer/mla.py - 48 lines added

Total: 171 lines added, 15 lines modified

---

Create Pull Request

Branch: claude/issue-7-20251207-2011 |