[TLE] Optimize fp8_fp4_paged_mqa_logits with a TLE TMA+WGMMA fast path

[henghengxiedaima]

PR Category

Operator

Type of Change

Performance Optimization

Description

Title: Optimize fp8_fp4_paged_mqa_logits with a TLE TMA+WGMMA fast path (SM90+)

Body:

Summary

Adds a deterministic TLE TMA + WGMMA fast path to the paged fp8_fp4_paged_mqa_logits operator. On Hopper (SM90+) with TLE available, shapes in the measured win band (D == 128, block_size == 256, max_ctx ∈ [12288, 16384]) are dispatched to a TMA-fed, shared-memory-resident WGMMA kernel; everything else keeps the existing Triton baseline. The change is purely additive — no existing code path is modified.

What was optimized

1. TLE TMA + WGMMA kernel (_mqa_logits_kernel_tle)

• Each KV page is bulk-copied directly from global memory into SMEM by the TMA engine (tle.gpu.copy on a TensorDescriptor over the fp8 paged KV cache [num_phys_blocks, block_size, head_dim]), synchronized via an mbarrier (tle.gpu.alloc_barrier / tle.gpu.barrier_wait). No per-lane register round-trip for KV.
• Q is loaded into SMEM once per CTA and reused across all pages (WGMMA operand A).
• The per-page attention is issued with tle.gpu.wgmma, reading both operands directly from SMEM.
• The page loop is software-pipelined with tl.range(num_stages=2), so the next page's TMA fill overlaps the current page's WGMMA + epilogue.
• Block-table lookups are safety-clamped so partial trailing tiles never read out of bounds; garbage pages are masked out of the store.

2. Deterministic dispatch guard (_can_use_tle)

• Fires only on SM90+ (Hopper) with triton >= 3.6 + TLE; master switch FLAGGEMS_FP8_FP4_PAGED_MQA_LOGITS_TLE (default on).
• Fires only when the KV page is a 32 KB fp8 tile (block_size == 256, D == 128) and the page loop is long enough for pipelining to reach steady state (max_ctx ∈ [12288, 16384], i.e. 6–8 pages per CTA at the operator's BLOCK_KV = 2048).
• Reuses the operator's existing _select_block_kv geometry, so the fast path is bit-identical to the baseline wherever it fires; all other shapes fall back to the unchanged baseline.

3. Performance (H800, Triton 3.6.0, vs the operator's Triton baseline)

• ≈1.2–1.3× in the enabled band (peak ≈1.3×, e.g. batch 128 / 32 at max_ctx = 16384).
• Output is bit-identical to the baseline where the fast path fires.

Issue

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}