Add the ORT-side FlashMLA launcher (still no call site)
Second of three steps. The previous commit got FlashMLA compiling; this one
adapts it to ORT's parameters and workspace conventions. Nothing calls it yet,
so behaviour is unchanged; wiring it into LaunchPagedLatentAttention behind an
env switch and measuring it in-model is the next commit.
flash_mla_latent.h deliberately does not include flash_mla.h. Those symbols only
exist when the sm90a object library is built, so the include stays inside the
.cu and callers see plain parameters.
Mapping notes, all cross-checked against the standalone benchmark
(dev/scripts/h200_3/flash_mla_bench.cu) that was already validated against an
fp32 reference:
- ORT's query heads become FlashMLA's `ngroups`, and `h` (its KV head count) is
1, so its flattened `seqlen_q` is seqlen_q * num_heads. ORT's query layout
(token, head, head_size) already matches its q_row_stride = head_size.
- k_ptr == v_ptr: the value is the leading 512 channels of the same 576-wide
latent row. This is the assumption that makes it MLA rather than ordinary
paged attention, and it is exactly what the ORT latent cache already stores.
- k_batch_stride is the PER-PAGE stride, not a per-sequence one, because the
cache is indexed through block_table.
- params.cu_seqlens_k is misnamed upstream: it is read as a per-sequence total
KV length, not a cumulative offset. ORT carries the length as of the start of
the step, so BuildSeqlensKKernel adds seqlen_q. It is a device kernel rather
than a host computation because the value must be right under CUDA graph
replay, where the host cannot see it.
- num_sm_parts follows the benchmark's formula (SMs / query tiles). Upstream's
own getFlashMlaNumSmParts ignores the KV length and returns the SM count (132
here) regardless, which is far too many for a short context.
FlashMlaLatentSupported() is deliberately strict. Every condition it screens --
bf16, unquantized cache, head_size 576, v_head_size 512, page size 64, one KV
head, no head sink, no kv_indices, no softcap, no sliding window, uniform query
length across the batch -- is something the kernel assumes rather than checks,
so a false positive would be silent corruption rather than a failure.
Verified to compile standalone at sm_90a with the same flags and include paths
the CUDA EP uses.
