Skip to content

[None][perf] Preallocate Kimi attention residual snapshots #17053

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
longlee0622 merged 3 commits into from
Aug 18, 2026
Merged

[None][perf] Preallocate Kimi attention residual snapshots #17053

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ad6c3b6
[perf] Preallocate Kimi attention residual snapshots
longlee0622 Jul 30, 2026
7d1d8b7
[test] Adapt snapshot parity test to general MLA
longlee0622 Aug 5, 2026
f7cc77c
[test] Remove misplaced Kimi snapshot model test
longlee0622 Aug 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 29 additions & 12 deletions tensorrt_llm/_torch/models/modeling_kimi_linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2281,26 +2281,30 @@ def forward(
self,
hidden_states: torch.Tensor,
block_residual: torch.Tensor,
num_snapshots: int,
attn_metadata: AttentionMetadata,
) -> Tuple[torch.Tensor, torch.Tensor]:
) -> Tuple[torch.Tensor, int]:
"""Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token).

Returns ``(prefix_sum, block_residual)`` with the snapshot stack in
kernel-native ``[K, M, H]`` layout; the running prefix sum is the
hidden state handed to the next layer.
``block_residual`` is a preallocated snapshot bank in kernel-native
``[K_max, M, H]`` layout. Returns the running prefix sum and the
number of valid bank rows.
"""
prefix_sum = hidden_states
valid_block_residual = block_residual[:num_snapshots]

if block_residual.shape[0] > 0:
if num_snapshots > 0:
hidden_states = _apply_attn_res(
prefix_sum,
block_residual,
valid_block_residual,
self.self_attention_res_proj,
self.self_attention_res_norm,
)

if self.layer_idx % self.attn_res_block_size == 0:
block_residual = torch.cat((block_residual, prefix_sum.unsqueeze(0)), dim=0)
block_residual[num_snapshots].copy_(prefix_sum)
num_snapshots += 1
valid_block_residual = block_residual[:num_snapshots]
prefix_sum = None

hidden_states = self.input_layernorm(hidden_states)
Expand All @@ -2312,7 +2316,7 @@ def forward(
prefix_sum = hidden_states

hidden_states = _apply_attn_res(
prefix_sum, block_residual, self.mlp_res_proj, self.mlp_res_norm
prefix_sum, valid_block_residual, self.mlp_res_proj, self.mlp_res_norm
)

hidden_states = self.post_attention_layernorm(hidden_states)
Expand All @@ -2324,7 +2328,7 @@ def forward(
hidden_states = self.mlp(hidden_states)

prefix_sum = prefix_sum + hidden_states
return prefix_sum, block_residual
return prefix_sum, num_snapshots


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -2360,6 +2364,9 @@ def __init__(self, model_config: ModelConfig):
cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype
)
self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype)
self.num_attn_res_snapshots = (
cfg.num_hidden_layers + cfg.attn_res_block_size - 1
) // cfg.attn_res_block_size

def forward(
self,
Expand All @@ -2384,9 +2391,16 @@ def forward(
"tokens); disable CUDA graphs and the overlap scheduler."
)

block_residual = hidden_states.new_zeros(0, hidden_states.shape[0], hidden_states.shape[1])
block_residual = hidden_states.new_empty(
self.num_attn_res_snapshots,
hidden_states.shape[0],
hidden_states.shape[1],
)
num_snapshots = 0
for layer in self.layers:
hidden_states, block_residual = layer(hidden_states, block_residual, attn_metadata)
hidden_states, num_snapshots = layer(
hidden_states, block_residual, num_snapshots, attn_metadata
)
if spec_metadata is not None:
# DFlash hidden-state capture. K3's attn-residual scheme
# already folds the residual into the running prefix sum
Expand All @@ -2398,7 +2412,10 @@ def forward(
spec_metadata.maybe_capture_hidden_states(layer.layer_idx, hidden_states, None)

hidden_states = _apply_attn_res(
hidden_states, block_residual, self.output_attn_res_proj, self.output_attn_res_norm
hidden_states,
block_residual[:num_snapshots],
self.output_attn_res_proj,
self.output_attn_res_norm,
)
return self.norm(hidden_states)

Expand Down
Loading