Filling the slot mappings with -1 is a correct approach to prevent dummy runs from corrupting the KV cache (specifically block 0). This ensures that the concat_and_cache kernels skip the write operation entirely.

However, this fix appears to be incomplete due to an existing issue in the surrounding code. At line 5364 (just above this hunk), _get_slot_mappings is called with num_tokens_padded=num_tokens, where num_tokens is the unpadded count. This means the returned sm tensors are only sized for the unpadded region. However, the model forward pass (and the underlying kernels) will use num_tokens_padded (the padded count from line 5344).

If num_tokens_padded > num_tokens, the kernels will read from the persistent slot_mapping buffer beyond the region filled by this new code. If that padding region contains zeros (which is likely for a persistent buffer), the corruption of block 0 will still occur for the padding tokens of the dummy run. To fully address this, line 5364 should be updated to use the local num_tokens_padded variable instead of num_tokens.

perhaps I was a bit hasty with the auto merge

I added this change num_tokens_padded=num_tokens_padded now as it seemed like a typo.

But the reason why this wasn't an issue was because:

            # Fill unused with -1. Needed for reshape_and_cache in full cuda
            # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID
            slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)

already handles the rest of the data.

Will run overnight test to see impact of num_tokens_padded=num_tokens_padded.