[Model][DSv4.1] FlashMLA mega attention and the NVFP4 compressed KV cache
The Issue & Context
Previous implementation of FlashMLA mega-attention and NVFP4 compressed KV cache was opt-in and based on an older branch; this PR rebases and supersedes that work, aligning it with recent KV-only context insertion changes.
Root Cause Analysis
The earlier approach required separate kernels and opt-in flags for query normalization, RoPE application, and KV cache handling, leading to complexity and inefficiency in fused attention computation.
Solution Strategy & Architectural Pattern
This PR introduces a unified FlashMLA mega-attention kernel that performs query RoPE, sparse attention, inverse RoPE on output, and FP8 casting in a single launch, reading and writing in a chunk-interleaved layout with runtime flags controlling query processing, and makes this kernel the default for SM100 hardware.
Key Takeaway for Contributors
Combining multiple attention operations into a single fused kernel with flexible query layout handling simplifies the codebase, improves performance, and sets a new default for optimized attention on SM100 GPUs.