vLLM Easier cudagraph integration

[LucasWilkinson]

This PR makes it easier to get a correct FlashMLA full cudagraph integration in vLLM where padding for cudagraph's happens after attention metadata data building (i.e. get_mla_metadata)

The changes

tile_scheduler_metadata0[0] = (now_idx >= batch_size ? -1 : now_idx);

and

if (begin_idx >= params.b || begin_idx < 0) return;

allows get_mla_metadata to be called with a smaller batch size than the graph was captured with since it now uses -1 to mark no-work tiles instead of params.b

The change

if (my_num_splits <= 1) {

makes it easier to pad out num_splits (cumulative) since we can just do:

self.cg_buf_num_splits[n:].fill_(num_splits[-1])

which pads with 0 split elements instead of have to pad with 1 split elements which is more complicated:

buffer_padding_size = self.cg_buf_num_splits.size(0) - n
self.cg_buf_num_splits[n:] = num_splits[-1] + 1 + torch.arange(buffer_padding_size)