[Performance]: DFlash / MLA DSpark decode spends ~0.2 ms per step in small bookkeeping kernels around the drafter #20548

Proposal to improve performance

With DFlash / DSpark speculative decoding and an MLA drafter (MLADSparkForCausalLM), each decode step runs a long chain of small torch kernels around the drafter, on the critical path:

• refreshing the draft block table (_refresh_ctx_block_tables);
• the gen-request bookkeeping in prepare_1st_drafter_inputs: bonus token, query / context positions, embedding ids, the noise block, the context-KV mask, the clamp and the paged write, the ctx_len advance and the advertised context length;
• in the drafter, the layer-invariant block fixup (_build_mla_block_fixup), a per-layer index_put of the block latents into the pool, and a contiguous copy plus an output copy per layer in the context-KV precompute.

Each of these is a handful of launches of a few microseconds.

Proposal: an opt-in switch under which each of these groups runs as one Triton kernel, bit-exact with the torch code it replaces. A second opt-in switch forks the gen tokens' context projection onto an aux stream, overlapping it with sample / accept. That projection reads only the captured target hidden states and the pre-step ctx_len.

Measured on Rubin, TP8 across 2 nodes (MNNVL), Kimi K3 with an MLA DSpark drafter, CUDA graphs, single request, decode step time with both switches: -0.195 / -0.206 / -0.186 / -0.212 / -0.224 / -0.210 ms at 2K / 32K / 64K / 96K / 128K / 393K context. Accepted tokens per step are unchanged.

PR to follow (draft).