[None][feat] AutoDeploy push the rope buffer to later stage - #13859
Conversation
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" --disable-fail-fast |
📝 WalkthroughWalkthroughThis PR refactors rotary positional embedding (RoPE) implementations across 13 auto_deploy custom model files by introducing a shared utility module ( ChangesUnified Rotary Embedding Refactoring
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_mistral3.py (1)
523-536:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftAvoid rebuilding full RoPE tables in every attention layer.
On Line 534 each
Mistral4Attentionlayer rebuilds full cos/sin tables frommax_position_embeddingsand then slices them. With large context configs, this is a significant hot-path regression (compute + memory) compared to the previous cached-read behavior.Please compute RoPE tables once per forward at model scope and pass sliced embeddings to layers, or otherwise cache/reuse per-request results across layers.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_mistral3.py` around lines 523 - 536, The attention forward currently calls self.rotary_emb(hidden_states, seq_len=seq_len) inside Mistral4Attention.forward, rebuilding full cos/sin tables per layer; fix by computing cos/sin once per model forward and passing the sliced embeddings into each layer (or add a simple cache inside rotary_emb keyed by seq_len/position_ids). Concretely: stop calling rotary_emb from Mistral4Attention.forward; update the model-level forward that runs the stack to call rotary_emb(...) once (using hidden_states/seq_len/position_ids) and pass the resulting cos and sin into each layer call (update Mistral4Attention.forward signature to accept cos, sin or add an argument like rope_cos, rope_sin), or alternatively modify rotary_emb to memoize the last generated tables and return slices without rebuilding; ensure usage sites reference the new parameters (forward, rotary_emb, position_ids, Mistral4Attention.forward).
🧹 Nitpick comments (2)
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_ir.py (1)
60-85: 💤 Low valueConsistency: add
attention_scalingattribute to match other refactored models.The function
build_rope_cos_sin_cachehas a defaultattention_scaling: float = 1.0, so this call works. However, other refactored models (Llama3RotaryEmbedding,DeepSeekV3RotaryEmbedding) explicitly define and passself.attention_scalingin their implementations for consistency.Consider adding
self.attention_scaling = 1.0in__init__and passing it:- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)) + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)) + self.attention_scaling = 1.0 self.register_buffer("inv_freq", inv_freq, persistent=False)And update the call:
- cos, sin = build_rope_cos_sin_cache(self.inv_freq, self.max_position_embeddings, x) + cos, sin = build_rope_cos_sin_cache(self.inv_freq, self.max_position_embeddings, x, self.attention_scaling)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_ir.py` around lines 60 - 85, The Qwen3RotaryEmbedding class should include an attention_scaling attribute for consistency with other rotary classes; add self.attention_scaling = 1.0 in Qwen3RotaryEmbedding.__init__ and update Qwen3RotaryEmbedding.forward to pass self.attention_scaling into build_rope_cos_sin_cache (i.e., call build_rope_cos_sin_cache(self.inv_freq, self.max_position_embeddings, x, self.attention_scaling)). This keeps behavior identical but matches the pattern used by Llama3RotaryEmbedding and DeepSeekV3RotaryEmbedding.tensorrt_llm/_torch/auto_deploy/models/custom/rotary_utils.py (1)
25-30: ⚡ Quick winAdd type annotations to
_apply.Please annotate
_applyarguments/return for mypy/lint consistency in this new module.♻️ Suggested patch
+from collections.abc import Callable + class RotaryEmbeddingBase(nn.Module): @@ - def _apply(self, fn): + def _apply(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> "RotaryEmbeddingBase": super()._apply(fn) inv_freq = getattr(self, "inv_freq", None) if isinstance(inv_freq, torch.Tensor) and inv_freq.is_floating_point(): self.inv_freq = inv_freq.float() return selfAs per coding guidelines "Always annotate functions; make the return type
Noneif the function does not return anything".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/auto_deploy/models/custom/rotary_utils.py` around lines 25 - 30, Annotate the _apply method signature: import Callable from typing and change def _apply(self, fn) to def _apply(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> torch.nn.Module (or the specific module class if preferred) so mypy knows the callable takes a Tensor and returns a Tensor and the method returns the module; keep the existing body (super()._apply(fn), inv_freq handling, return self) and ensure typing imports are added at top of the file.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_mistral3.py`:
- Around line 523-536: The attention forward currently calls
self.rotary_emb(hidden_states, seq_len=seq_len) inside
Mistral4Attention.forward, rebuilding full cos/sin tables per layer; fix by
computing cos/sin once per model forward and passing the sliced embeddings into
each layer (or add a simple cache inside rotary_emb keyed by
seq_len/position_ids). Concretely: stop calling rotary_emb from
Mistral4Attention.forward; update the model-level forward that runs the stack to
call rotary_emb(...) once (using hidden_states/seq_len/position_ids) and pass
the resulting cos and sin into each layer call (update Mistral4Attention.forward
signature to accept cos, sin or add an argument like rope_cos, rope_sin), or
alternatively modify rotary_emb to memoize the last generated tables and return
slices without rebuilding; ensure usage sites reference the new parameters
(forward, rotary_emb, position_ids, Mistral4Attention.forward).
---
Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_ir.py`:
- Around line 60-85: The Qwen3RotaryEmbedding class should include an
attention_scaling attribute for consistency with other rotary classes; add
self.attention_scaling = 1.0 in Qwen3RotaryEmbedding.__init__ and update
Qwen3RotaryEmbedding.forward to pass self.attention_scaling into
build_rope_cos_sin_cache (i.e., call build_rope_cos_sin_cache(self.inv_freq,
self.max_position_embeddings, x, self.attention_scaling)). This keeps behavior
identical but matches the pattern used by Llama3RotaryEmbedding and
DeepSeekV3RotaryEmbedding.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/rotary_utils.py`:
- Around line 25-30: Annotate the _apply method signature: import Callable from
typing and change def _apply(self, fn) to def _apply(self, fn:
Callable[[torch.Tensor], torch.Tensor]) -> torch.nn.Module (or the specific
module class if preferred) so mypy knows the callable takes a Tensor and returns
a Tensor and the method returns the module; keep the existing body
(super()._apply(fn), inv_freq handling, return self) and ensure typing imports
are added at top of the file.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: de7f6d9f-b0a5-4c3c-a42e-824a38de7161
📒 Files selected for processing (14)
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_ir.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma3n.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma4.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_kimi_k2.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_llama3_ir.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_llama4.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_minimax_m2.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_mistral3.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_ir.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_moe.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_starcoder2.pytensorrt_llm/_torch/auto_deploy/models/custom/rotary_utils.py
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" --disable-fail-fast |
1 similar comment
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" --disable-fail-fast |
|
PR_Github #47241 [ run ] triggered by Bot. Commit: |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
PR_Github #47241 [ run ] completed with state
|
govind-ramnarayan
left a comment
govind-ramnarayan
left a comment
There was a problem hiding this comment.
LGTM, does it have any perf implications?
|
Can we try running these locally on top of #13925 The following are broken and waived after transformers upgrade: It'll be safe to test these as well. Thanks! |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma3n.py
@govind-ramnarayan Should have no perf implications as the optimize_rope (a later transform) will do the buffer init... So ideally it will be functionally same.. |
bmarimuthu-nv
left a comment
bmarimuthu-nv
left a comment
There was a problem hiding this comment.
Thanks @nvchenghaoz , adding a transform unit test would be helpful!
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py # tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py # tensorrt_llm/_torch/auto_deploy/models/custom/modeling_minimax_m2.py
|
/bot run --stage-list "A10-Build_Docs, A10-PackageSanityCheck-PY310-UB2204, A100X-PackageSanityCheck-PY312-UB2404, A30-AutoDeploy-1, H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, A100X-PyTorch-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" --disable-fail-fast |
|
PR_Github #49458 [ run ] triggered by Bot. Commit: |
|
PR_Github #49458 [ run ] completed with state
|
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" --disable-fail-fast |
|
PR_Github #49548 [ run ] triggered by Bot. Commit: |
|
PR_Github #49548 [ run ] completed with state
|
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" --disable-fail-fast |
|
PR_Github #49731 [ run ] triggered by Bot. Commit: |
|
PR_Github #49731 [ run ] completed with state
|
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1, DGX_B200-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1, DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1" --disable-fail-fast |
|
PR_Github #50005 [ run ] triggered by Bot. Commit: |
|
PR_Github #50005 [ run ] completed with state |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_ir.py # tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3.py
|
/bot run --disable-fail-fast |
|
PR_Github #50367 [ run ] triggered by Bot. Commit: |
|
PR_Github #50367 [ run ] completed with state
|
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #50391 [ run ] triggered by Bot. Commit: |
|
PR_Github #50391 [ run ] completed with state |
…3859) Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
5 participants
The main motivation for this PR is to unblock the pipeline cache work. As of now, the pipeline cache design will cache all the buffers. That means the buffer init in the modeling file will be saved into the cache directory. The issue here is the size of the rope is large (could be 500MB), and it will be saved for all ranks. The total size of the rope buffer in the cache will be Rope_size * world_size..
To mitigate this issue, push the rope buffer pre-calculation out of the modeling file.
Summary by CodeRabbit
Release Notes