Accelerating Mixture of Block Attention via Dynamic Routing, KV-Centric Execution and Runtime Pipelining

Mixture of Block Attention (MoBA) reduces long-context computation by selecting key-value (KV) blocks for each query. FlashMoBA implements MoBA efficiently but retains a fixed Top-$k$ budget and query-centric schedule. Fixed Top-$k$ applies one budget across different score distributions and couples routing to task granularity; query-range ownership fragments K/V reuse when several ranges select the same column. We present LightningMoBA, which co-designs selection and execution through dynamic routing, KV-centric ownership, and runtime pipelining. Dynamic routing thresholds MoBA's per-query block scores and directly emits variable-cardinality sparse metadata. KV-centric ownership groups selected rows by (query head, KV block) column, splits high-fanout columns into bounded tasks, and reuses each task's loaded KV tile. On Hopper, pipelining overlaps regular K/V transfers and irregular query gathers with Tensor Core computation. Together these mechanisms decouple retention from GPU scheduling granularity. At matched Top-$16$ retention across five checkpoints, LightningMoBA accelerates end-to-end generation by $1.20$ to $1.60\times$ while raising mean RULER score from $74.42$ to $75.28$. In controlled two-model replays, complete device-path speedup grows from $1.56$ to $1.65\times$ at Top-$16$ to $2.16$ to $2.56\times$ at Top-$32$, while routing and CSC construction reach $6.99$ to $11.25\times$ at Top-$64$. The code of LightningMoBA is available at https://anonymous.4open.science/r/LightningMoba-A090