[Relax][Frontend][KVCache] Add masked sequence prefill helper for encoder valid lengths

[xthomaswang]

Summary

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.

Tests

tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py covers:

• valid_len == 0
• valid_len == seq_len
• mixed valid lengths within a batch
• grouped-query attention (h_q > h_kv)

Run with:

python -m pytest -v tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py

[gemini-code-assist]

Code Review

This pull request introduces a new masked sequence prefill kernel, _attention_sequence_prefill_with_mask, designed for batched encoder-style inputs with right-padding. It includes a corresponding test suite covering various edge cases like zero-length sequences and grouped-query attention. Feedback includes removing an unused causal parameter, cleaning up unnecessary pylint suppressions for the lse buffer, and adding explicit bounds checks for kv_len during K and V loading to ensure robustness.

[xthomaswang]

Addressed the suggestions from gemini