Performance comparison for dense shared prefix + block-diagonal local mask: custom mask in BatchPrefillWithPagedKVCacheWrapper vs MultiLevelCascadeAttentionWrapper

[catnanami]

I am trying to evaluate the best way to implement an attention pattern with a dense shared/global prefix plus a local block-diagonal suffix.

The mask pattern is shown below:

I am considering two implementations:

BatchPrefillWithPagedKVCacheWrapper with a custom mask
MultiLevelCascadeAttentionWrapper

Questions:

1. For this pattern, is cascade attention expected to be faster than paged prefill with a custom mask?
2. Or is this pattern better handled by a custom kernel/backend?

My use case is speculative decoding, where all verification tokens share the same prefix KV, but each local block should only attend to its own block.

[VectorPeak]

For the mask in the image - a dense shared prefix on the left plus independent block-diagonal local regions on the right - MultiLevelCascadeAttentionWrapper is usually the better fit than BatchPrefillWithPagedKVCacheWrapper with one large custom mask.

The reason is that the custom-mask prefill path does not really exploit this structure as well: it still treats the problem like one masked attention problem over the full layout, so the block-diagonal sparsity mostly becomes masking overhead. Cascade attention lets you model this as:

• shared/global prefix KV: attended by all verification tokens
• local suffix KV: split into independent blocks/requests
• final result: merge prefix-attention state with each local block’s attention state

For speculative decoding, this is exactly the kind of case where cascade can help, especially when the shared prefix is large and the local verification blocks are relatively small. Each local block can be represented as a separate batch item sharing the same prefix. If the local part also needs causal masking inside each block, use the normal causal/local mask for that per-block suffix attention.