[None][perf] Optimize MiniMax-M3 CTX block selection

[peihu-nv]

@coderabbitai summary

Description

MiniMax-M3 sparse planning had two independent costs in pure prefill: the
general block-selector kernel was inefficient for the common at-most-128-block
shape, and its Q-major output was materialized into a separate head-major
tensor immediately before K2Q planning.

This PR makes two complementary optimizations:

1. Specialize block selection for at most 128 allocated blocks. Dedicated
   warp kernels cover at most 32, 64, and 128 candidates using register-resident
   selection and sorting. Shapes above 128 continue to use the existing general
   selector.
2. Make the selector-to-K2Q handoff zero-copy in pure prefill. The selector
   writes directly into head-major backing storage while preserving the public
   logical [Q, Hkv, 16] view. The existing
   result.permute(1, 0, 2).contiguous() K2Q handoff therefore becomes a
   contiguous view instead of launching a materialization copy. This layout
   optimization also applies when shapes above 128 use the general selector;
   mixed and decode batches retain the existing Q-major backing.

In a matched GB200 1P1D run (TP4/EP4 per server, 8K input / 1K output,
concurrency 64):

• The selector kernel improved from 6.264 to 3.951 ms per 8K prefill
  (-36.9%).
• The zero-copy handoff removed 57 selector-to-K2Q materializations per
  prefill; selected int32 copy work fell from 61 calls / 0.181 ms to 4 calls /
  0.008 ms.
• Combined sparse planning improved from 11.230 to 8.818 ms (-21.5%),
  resulting in approximately +1.1% context throughput.

Test Coverage

• Focused selector correctness suite: 47 passed, covering random inputs,
  ties, forced blocks, 64/128-block dispatch, output strides, and CUDA graph
  replay.
• Pure-prefill/mixed/decode routing suite: 3 passed.
• Long-query handoff checks at Q=6599, 7379, and 8175 verified shared storage,
  contiguous head-major K2Q input, and zero CUDA events during
  permute(...).contiguous().
• SM100 wheel build job 2469230 completed; corrected validation job
  2469447 completed.
• Matched E2E/Nsys job 2469565 completed 1,000/1,000 requests. Normalized
  8K results: selector 6.264 to 3.951 ms (-36.9%), sparse planning 11.230 to
  8.818 ms (-21.5%), and CTX forward 318.359 to 314.991 ms (-1.06%),
  equivalent to approximately +1.1% context throughput.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of up-to-date bot commands, post /bot help as a PR comment and check the bot's reply.