Adds `_attention_sequence_prefill_with_mask` in
`python/tvm/relax/frontend/nn/llm/kv_cache.py` — a masked variant of the
existing sequence prefill kernel that supports right-padded encoder
batches with per-sample `valid_lens`.
The existing `_attention_sequence_prefill` assumes all positions in `[0,
seq_len)` are valid, which breaks for padded encoder inputs where each
batch element has a different valid prefix length. This helper adds the
masking semantics needed for correctness:
- accepts a per-batch `valid_lens` input
- ignores padded query rows and padded key/value positions
- excludes padded `(row, col)` pairs from the online softmax update
It reuses the existing prefill kernel config and schedule — no new
tuning knobs, no target-specific changes, no performance claims.
Correctness only.
## Motivation: encoder batch prefill for downstream consumers
This is the TVM-side primitive needed to support **encoder batch
prefill** in downstream projects like `mlc-llm`, where padded encoder
batches with `valid_lens` need to be lowered without materializing an
explicit broadcast attention mask on the host.
The helper is generic and useful for any encoder-style sequence prefill
consumer with per-sample valid lengths.
Xijing Wang
a6e2ea8ac8