[perf] Use direct copy (broadcast) instead of cat for k_nope/k_pe in MLA prefill

[minosfuture]

Purpose

Reduce k tensor concatenation latency, from 3.16ms to 1.61ms for batch size 32768 (i.e., k.shape=torch.Size([32768, 128, 192]), k_nope.shape=torch.Size([32768, 128, 128]), k_pe.shape=torch.Size([32768, 1, 64]).

Test Plan

• verified perf improvement with trace
• verified accuracy

Test Result

local-completions (model=nvidia/DeepSeek-R1-0528-FP4,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=32), gen_kwargs: (None), limit: 100.0, num_fewshot: None, batch_size: 1

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.96	±	0.0197
		strict-match	5	exact_match	↑	0.96	±	0.0197

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

This pull request introduces a performance optimization for MLA prefill by replacing torch.cat with a more efficient direct copy and broadcast mechanism in a new _concat_k_nope_k_pe method. This change correctly avoids the overhead of concatenating expanded, non-contiguous tensors and is applied consistently in _compute_prefill_context, _context_parallel_compute_prefill_context, and _forward_prefill. The implementation is clean and well-documented, leading to the significant latency reduction shown in the PR description. The change looks good.