[vsabavat]

Description

trtllm::kda_prefill (tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py) returned a view of a buffer
it reuses, while its fake declares a new tensor.

• What goes wrong: _get_buffers caches the prefill's intermediate tensors per input shape (_buf_cache), and
  the output was one of them (O_flat). Every call with the same shapes wrote the same storage, so a caller's output
  changed under it at the next call with those shapes. The op's fake returns v.new_empty(v.shape), a new tensor
  that aliases nothing, so a compiler tracing through the op has no reason to finish reading one call's output before
  the next call runs.
• Before: on main, a second call with the same shapes and other inputs rewrote the first call's output, on the
  varlen path and on the padded equal-length path.
• Who is affected: in-tree, the KDA mixer (modules/kimi_kda/kimi_kda_mixer.py) consumes the output in the same
  layer (_output_gate), before the next layer's call, so eager runs were correct. Any caller that keeps an output
  across calls, or a compiled graph that orders work by the fake's contract, reads the wrong values.
• Fix: the output is allocated on every call, and K4 writes it directly; the cache no longer holds it. On the
  varlen path it keeps one chunk of slack past the last token, like the cached buffers. The op now returns a tensor
  of its own, as the fake declares.
• Cost (GB200, Kimi K3's KDA shapes: 6 and 24 heads of 128, 1k-8k tokens, one to four sequences):
  • no GPU time;
  • about 35 us more host time per call, which shows where a call is host-bound (up to 4k tokens at 6 heads: ~235
    to ~270 us per call).
  • Cloning the cached output at the op boundary instead costs about the same host time, plus a copy: 5 us at
    6 heads x 8k tokens, 15 us at 24 heads.

Test Coverage

tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py::test_kda_prefill_op_output_survives_the_next_call
(new; varlen [100, 200] and one 300-token row without cu_seqlens): keeps the first call's output, calls again
with the same shapes and other inputs, and checks the first output is unchanged and the two differ. On GB200 both
cases fail on main and pass with this PR. test_kda_prefill_op.py and the other KDA prefill tests
(modules/kimi_kda/, modeling/test_kimi_kda_fp8_packed_prefill.py) have no new failures. CI collects the file
through l0_b200.yml (pre-merge, single GPU).

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. (No API change.)

• Any new dependencies have been scanned for license and vulnerabilities (No new dependencies.)

• CODEOWNERS updated if ownership changes (No ownership change.)

• Documentation updated as needed (No user-facing change.)

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

• The reviewers assigned automatically/manually are appropriate for the PR. (To check once the PR is open.)

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

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[vsabavat]

Superseded by the consolidated draft PR: #19988. Source commits and follow-ups are retained; prior reviews remain historical evidence, not approval of the replacement heads.