[#11032][None] AutoDeploy support for GLM 4.7 flash - #11150
Closed
bmarimuthu-nv wants to merge 20 commits into
Closed
[#11032][None] AutoDeploy support for GLM 4.7 flash#11150bmarimuthu-nv wants to merge 20 commits into
bmarimuthu-nv wants to merge 20 commits into
Conversation
lucaslie
reviewed
Jan 30, 2026
lucaslie
reviewed
Feb 3, 2026
Contributor
There was a problem hiding this comment.
this should be configured as --args.yaml-extra not as --yaml-extra
1 task
lucaslie
reviewed
Feb 3, 2026
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Collaborator
|
@bmarimuthu-nv : is this ready for review? can it be moved out of draft and cleaned up? |
bmarimuthu-nv
marked this pull request as ready for review
February 4, 2026 21:12
bmarimuthu-nv
requested review from
FrankD412,
MrGeva,
laikhtewari and
nv-guomingz
February 4, 2026 21:12
suyoggupta
changed the title
Contributor
📝 WalkthroughWalkthroughThis PR restructures MLA (Multi-Head Latent Attention) from a single monolithic implementation into a modular backend architecture with multiple execution paths (torch, torch_backend, flashinfer). It introduces new model implementations for DeepSeekV3 and GLM4 MoE Lite, removes the patching mechanism, expands logging/debugging utilities, and adds comprehensive test coverage for the new backends. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
🚥 Pre-merge checks | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 14
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py (1)
1-7:⚠️ Potential issue | 🟠 MajorAdd the standard NVIDIA copyright header.
This source file is missing the required NVIDIA copyright header with the latest modification year. Please add the repo’s standard header before the module docstring.
tensorrt_llm/_torch/auto_deploy/shim/interface.py (1)
323-330:⚠️ Potential issue | 🟠 MajorGuard
CachedSequenceInterface.to()against unallocated contiguous caches.
Contiguous paged resources are stored asNoneuntil after KVCacheManager creation. If.to()is called earlier, it will attemptNone.to(...)and raise. Also, the currentto()path discards the returned tensor.🛠️ Suggested fix
def to(self, *args, **kwargs) -> None: self.info.to(*args, **kwargs) # Only move locally-allocated caches (paged/state caches are managed by cache managers) for name, cache in self._caches.items(): - if name not in self._paged_cache_order and name not in self._state_resource_order: - cache.to(*args, **kwargs) + if ( + name not in self._paged_cache_order + and name not in self._state_resource_order + and name not in self._contiguous_paged_cache_order + ): + if cache is not None: + self._caches[name] = cache.to(*args, **kwargs)
🤖 Fix all issues with AI agents
In `@docs/source/torch/auto_deploy/advanced/workflow.md`:
- Around line 18-20: The example config uses "insert_cached_mla_attention":
{"backend": "torch_mla"} while the current default is "flashinfer_mla"; update
the example or add a clarifying note: either change "torch_mla" to
"flashinfer_mla" for consistency, or add a comment next to
insert_cached_mla_attention (and optionally insert_cached_attention) explaining
both supported backends ("torch_mla" vs "flashinfer_mla") and when to pick each
one, and ensure compile_model/backends still align with the chosen MLA backend.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py`:
- Around line 1-13: Add the missing NVIDIA copyright header to the top of the
module (above the module docstring) following the project's standard header
format (include organization name, year range, and any SPDX identifier if used).
Update tensorrt_llm._torch.auto_deploy.custom_ops.mla.__init__ (the file that
exports MultiHeadLatentAttention, FlashInferMLAAttention, torch_mla,
torch_backend_mla_with_cache, flashinfer_mla_with_cache) by inserting the
canonical NVIDIA header block as the very first lines of the file so the
existing docstring and imports remain unchanged.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py`:
- Line 1: Add the required NVIDIA copyright header to the top of
flashinfer_mla.py: insert the standard NVIDIA file header comment (including
"Copyright (c) <year>, NVIDIA CORPORATION. All rights reserved." with the year
being the latest meaningful modification) before the module docstring in
tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py so the
file-level header precedes the existing triple-quoted module docstring and
complies with project coding guidelines.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py`:
- Line 1: The file lacks the required NVIDIA copyright header; add the standard
NVIDIA copyright comment block (including the year of latest meaningful
modification) at the very top of
tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py before the
existing module docstring ("""Custom ops for MultiHead Latent Attention (MLA)
with FlashInfer-compatible cache."""), ensuring the header follows the project's
canonical format and includes the correct year.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py`:
- Line 1: The file torch_mla.py is missing the required NVIDIA copyright header;
prepend the standard NVIDIA copyright header block (including "Copyright (c)
<year>, NVIDIA CORPORATION. All rights reserved." with the year set to the
latest meaningful modification) at the top of torch_mla.py before the module
docstring ("""Torch reference implementation for Multi-head Latent Attention
(MLA).""") so the file complies with the TensorRT-LLM source file guidelines.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py`:
- Around line 26-47: The function _update_kv_cache mutates k_cache and v_cache
in-place but is annotated to return torch.Tensor; change the signature to return
None (-> None) to reflect in-place behavior, and ensure there are no return
statements in _update_kv_cache so type checkers and callers treat it as a void
function.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py`:
- Around line 1-4: Add the required NVIDIA copyright header at the top of this
source file (above the imports) so the file that defines/exports
DeepSeekV3ForCausalLM, Glm4MoeLiteForCausalLM, NemotronFlashForCausalLM,
NemotronFlashPreTrainedTokenizerFast, and NemotronHForCausalLM includes the
company header text; insert the standard multi-line header block exactly as used
in other project files and ensure it precedes all code and imports.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py`:
- Line 1: Update the file header in modeling_deepseek.py to reflect the correct
year: change the copyright line that currently says "Copyright (c) 2025, NVIDIA
CORPORATION. All rights reserved." to use 2026 instead; ensure the updated
header remains otherwise unchanged and matches the format used across other
TensorRT-LLM source files.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py`:
- Around line 1-2: Update the file header year from 2025 to 2026 by editing the
top copyright comment; specifically change the line beginning with "Copyright
(c) 2025, NVIDIA CORPORATION. All rights reserved." to use 2026 so the file
(modeling_glm4_moe_lite.py) reflects the latest modification year.
In `@tensorrt_llm/_torch/auto_deploy/shim/interface.py`:
- Around line 346-352: Contiguous paged caches in _contiguous_paged_cache_order
are left at the old shape after resize_kv_cache_manager() recreates
_kv_cache_manager, causing shape mismatches; update the code so that after
_kv_cache_manager is resized/recreated you recompute num_pages =
self._kv_cache_manager.blocks_in_primary_pool and page_size =
self._kv_cache_manager.tokens_per_block and then reallocate each contiguous
cache by calling handler.allocate(self.info, num_pages, page_size) and storing
the result back into self._caches[name] (optionally releasing or overwriting the
old cache first) to ensure caches match the new KVCacheManager shape.
In `@tensorrt_llm/_torch/auto_deploy/transform/optimizer.py`:
- Around line 73-77: Remove the commented-out debug block (the lines referencing
t_name == "compile_model", run_comparison, and breakpoint) from
tensorrt_llm/_torch/auto_deploy/transform/optimizer.py to eliminate debugging
artifacts and fix the formatting error; if you want to preserve the debug
capability instead of deleting it, wrap the logic in a conditional gated by an
environment variable (similar to AD_DUMP_GRAPHS_DIR), e.g., check os.environ and
only call run_comparison(mod, cm, self.factory, output_dir=...) when enabled,
ensuring no leftover commented code or stray breakpoint() calls remain.
In `@tensorrt_llm/_torch/auto_deploy/utils/logger.py`:
- Around line 1-11: The project imports dill and safetensors in
tensorrt_llm._torch.auto_deploy.utils.logger (the symbols dill and
safetensors.torch.save_file / safetensors_save), but these packages aren't
declared as install-time dependencies; update the project's dependency metadata
(e.g., requirements.txt, setup.py install_requires, or pyproject.toml
[project.dependencies]) to include dill and safetensors (use the appropriate
package names on PyPI) and pin or specify compatible version ranges, then run
dependency install to validate imports succeed.
In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py`:
- Around line 359-392: The disable_overlap_scheduler flag in
TestGLM4Flash.get_default_kwargs is set to False but differs from the model
registry's configuration; update get_default_kwargs to match the registry flag
(set "disable_overlap_scheduler": True) to keep behavior consistent, or
explicitly document the intentional mismatch by adding a clear comment in the
TestGLM4Flash class and in get_default_kwargs noting why tests require False
while the deployed model uses True; refer to the TestGLM4Flash class and
get_default_kwargs to locate the change and to the "disable_overlap_scheduler"
key to update or document.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_glm4_moe_lite_modeling.py`:
- Around line 485-523: Remove the debug print statements in the numerical
equivalence test: delete or gate all prints that reference hf_state_dict,
custom_state_dict, custom_moe, hf_gate_up, converted_gate_0, and loaded_gate_0
(the blocks printing shapes, slices, and the "=== ..." headers). Keep the actual
assertions and the call to _convert_hf_moe_state_dict_to_custom and
custom_moe.load_state_dict as-is; if logging is desired for debugging, wrap the
prints behind an explicit debug flag (e.g., a module-level DEBUG or an env var)
so CI runs remain quiet.
🧹 Nitpick comments (9)
examples/auto_deploy/deepseek_mini.yaml (1)
1-2: Clarify the config path in the run command.If users run from repo root, the current comment may be ambiguous. Consider using the explicit relative path for the YAML.
💡 Suggested comment tweak
-# python build_and_run_ad.py --model "deepseek-ai/DeepSeek-R1" --args.yaml-extra "deepseek_mini.yaml" +# python build_and_run_ad.py --model "deepseek-ai/DeepSeek-R1" --args.yaml-extra "examples/auto_deploy/deepseek_mini.yaml"tensorrt_llm/_torch/auto_deploy/utils/logger.py (1)
317-353: Broad exception handling is acceptable for debug utilities, but consider logging exception types.The
except Exceptionclauses (Lines 321, 333, 352) are flagged by static analysis, but in debug/dump utilities, this pattern is reasonable—failing to dump artifacts shouldn't crash the main pipeline.However, logging only
{e}may lose valuable context. Consider including the exception type for easier debugging:💡 Optional: Include exception type in warning messages
except Exception as e: - self.warning(f"Failed to dump GraphModule: {e}") + self.warning(f"Failed to dump GraphModule: {type(e).__name__}: {e}")Apply similarly to lines 333-334 and 352-353.
tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py (1)
1028-1067: AlignContiguousPagedResourceHandler.allocatewith theResourceHandlercontract.
This override adds required parameters (num_pages,page_size), so any code treating handlers polymorphically asResourceHandlercould callallocate(sequence_info)and crash. Consider keepingallocate(sequence_info)as the contract and exposing a separate method (or optional kwargs) for the contiguous path.tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py (1)
19-30: Prefer module imports to preserve namespace.
This file imports individual symbols from several modules (e.g.,from transformers import AutoConfig, PretrainedConfig). The codebase guideline asks to keep module namespaces instead (e.g.,import transformersand usetransformers.AutoConfig).As per coding guidelines: Always maintain the namespace when importing Python modules, even if only one class or function from a module is used.
tensorrt_llm/bench/benchmark/__init__.py (1)
71-75: Remove the unusednoqa(RUF100).
Ruff flags the# noqa: F401as unused because F401 isn’t enabled. Either drop it or enable F401 if you want to enforce unused‑import checks.🧹 Suggested cleanup
- import tensorrt_llm._torch.auto_deploy.models.custom # noqa: F401 + import tensorrt_llm._torch.auto_deploy.models.customtensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py (1)
15-21: Consider sorting__all__alphabetically.Ruff (RUF022) suggests applying isort-style sorting to
__all__. This is a minor style improvement.Suggested fix
__all__ = [ + "FlashInferMLAAttention", "MultiHeadLatentAttention", - "FlashInferMLAAttention", - "torch_mla", - "torch_backend_mla_with_cache", "flashinfer_mla_with_cache", + "torch_backend_mla_with_cache", + "torch_mla", ]examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml (1)
11-12: Remove duplicate commentedattn_page_sizeentry.There are two commented
attn_page_sizelines with different values (512 and 64). This appears to be leftover from experimentation. Consider removing one to avoid confusion.Suggested fix
# max_seq_len: 512 - # attn_page_size: 512 # attn_page_size: 64 # attn_backend: flashinferexamples/auto_deploy/cookbooks/glm_4.7_flash_trtllm_cookbook.ipynb (1)
151-158: Prefer--configover--extra_llm_api_options.Per coding guidelines, prefer using
--configover--extra_llm_api_optionswhen documenting CLI commands for TensorRT-LLM tools.Suggested fix
"trtllm-serve \"zai-org/GLM-4.7-Flash\" \\\n", " --host 0.0.0.0 \\\n", " --port 8000 \\\n", " --backend _autodeploy \\\n", " --trust_remote_code \\\n", -" --extra_llm_api_options glm4.7_flash.yaml\n", +" --config glm4.7_flash.yaml\n",tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_mla_op.py (1)
192-192: Add explicitstrict=parameter tozip().Per Ruff B905, using
zip()without an explicitstrict=parameter can mask length mismatches. Sinceinput_positionsandseq_lengthsshould have the same length, usestrict=True.Suggested fix
- kv_lengths = [pos + seq_len for pos, seq_len in zip(input_positions, seq_lengths)] + kv_lengths = [pos + seq_len for pos, seq_len in zip(input_positions, seq_lengths, strict=True)]
Comment on lines
18
to
20
| "insert_cached_attention": {"backend": "flashinfer"}, # or "triton" | ||
| "insert_cached_mla_attention": {"backend": "MultiHeadLatentAttention"}, | ||
| "insert_cached_mla_attention": {"backend": "torch_mla"}, | ||
| "compile_model": {"backend": "torch-compile"}, |
Contributor
There was a problem hiding this comment.
Clarify backend choice in the example.
The default config now sets flashinfer_mla, while this example uses torch_mla. Consider noting both options (or aligning to the default) to reduce confusion for users following the workflow.
🤖 Prompt for AI Agents
In `@docs/source/torch/auto_deploy/advanced/workflow.md` around lines 18 - 20, The
example config uses "insert_cached_mla_attention": {"backend": "torch_mla"}
while the current default is "flashinfer_mla"; update the example or add a
clarifying note: either change "torch_mla" to "flashinfer_mla" for consistency,
or add a comment next to insert_cached_mla_attention (and optionally
insert_cached_attention) explaining both supported backends ("torch_mla" vs
"flashinfer_mla") and when to pick each one, and ensure compile_model/backends
still align with the chosen MLA backend.
Comment on lines
+1
to
+13
| """MLA (Multi-head Latent Attention) custom ops. | ||
|
|
||
| Exports: | ||
| - MultiHeadLatentAttention: Attention descriptor for MLA (registered as "torch_mla") | ||
| - FlashInferMLAAttention: Attention descriptor for FlashInfer MLA (registered as "flashinfer_mla") | ||
| - torch_mla: Source op for MLA attention | ||
| - torch_backend_mla_with_cache: Cached backend op with FlashInfer-compatible cache | ||
| - flashinfer_mla_with_cache: Cached backend op using FlashInfer MLA kernels | ||
| """ | ||
|
|
||
| from .flashinfer_mla import FlashInferMLAAttention, flashinfer_mla_with_cache | ||
| from .torch_backend_mla import MultiHeadLatentAttention, torch_backend_mla_with_cache | ||
| from .torch_mla import torch_mla |
Contributor
There was a problem hiding this comment.
Missing NVIDIA copyright header.
This source file is missing the required NVIDIA copyright header.
Suggested fix
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
"""MLA (Multi-head Latent Attention) custom ops.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| """MLA (Multi-head Latent Attention) custom ops. | |
| Exports: | |
| - MultiHeadLatentAttention: Attention descriptor for MLA (registered as "torch_mla") | |
| - FlashInferMLAAttention: Attention descriptor for FlashInfer MLA (registered as "flashinfer_mla") | |
| - torch_mla: Source op for MLA attention | |
| - torch_backend_mla_with_cache: Cached backend op with FlashInfer-compatible cache | |
| - flashinfer_mla_with_cache: Cached backend op using FlashInfer MLA kernels | |
| """ | |
| from .flashinfer_mla import FlashInferMLAAttention, flashinfer_mla_with_cache | |
| from .torch_backend_mla import MultiHeadLatentAttention, torch_backend_mla_with_cache | |
| from .torch_mla import torch_mla | |
| # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """MLA (Multi-head Latent Attention) custom ops. | |
| Exports: | |
| - MultiHeadLatentAttention: Attention descriptor for MLA (registered as "torch_mla") | |
| - FlashInferMLAAttention: Attention descriptor for FlashInfer MLA (registered as "flashinfer_mla") | |
| - torch_mla: Source op for MLA attention | |
| - torch_backend_mla_with_cache: Cached backend op with FlashInfer-compatible cache | |
| - flashinfer_mla_with_cache: Cached backend op using FlashInfer MLA kernels | |
| """ | |
| from .flashinfer_mla import FlashInferMLAAttention, flashinfer_mla_with_cache | |
| from .torch_backend_mla import MultiHeadLatentAttention, torch_backend_mla_with_cache | |
| from .torch_mla import torch_mla |
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py` around lines 1 -
13, Add the missing NVIDIA copyright header to the top of the module (above the
module docstring) following the project's standard header format (include
organization name, year range, and any SPDX identifier if used). Update
tensorrt_llm._torch.auto_deploy.custom_ops.mla.__init__ (the file that exports
MultiHeadLatentAttention, FlashInferMLAAttention, torch_mla,
torch_backend_mla_with_cache, flashinfer_mla_with_cache) by inserting the
canonical NVIDIA header block as the very first lines of the file so the
existing docstring and imports remain unchanged.
| @@ -0,0 +1,954 @@ | |||
| """FlashInfer-based MLA (Multi-head Latent Attention) backend with paged caching. | |||
Contributor
There was a problem hiding this comment.
Add the required NVIDIA copyright header.
🔧 Proposed fix
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
"""FlashInfer-based MLA (Multi-head Latent Attention) backend with paged caching.As per coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of latest meaningful modification.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| """FlashInfer-based MLA (Multi-head Latent Attention) backend with paged caching. | |
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | |
| """FlashInfer-based MLA (Multi-head Latent Attention) backend with paged caching. |
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` at line 1,
Add the required NVIDIA copyright header to the top of flashinfer_mla.py: insert
the standard NVIDIA file header comment (including "Copyright (c) <year>, NVIDIA
CORPORATION. All rights reserved." with the year being the latest meaningful
modification) before the module docstring in
tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py so the
file-level header precedes the existing triple-quoted module docstring and
complies with project coding guidelines.
| @@ -0,0 +1,518 @@ | |||
| """Custom ops for MultiHead Latent Attention (MLA) with FlashInfer-compatible cache. | |||
Contributor
There was a problem hiding this comment.
Add the required NVIDIA copyright header.
🔧 Proposed fix
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
"""Custom ops for MultiHead Latent Attention (MLA) with FlashInfer-compatible cache.As per coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of latest meaningful modification.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| """Custom ops for MultiHead Latent Attention (MLA) with FlashInfer-compatible cache. | |
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | |
| """Custom ops for MultiHead Latent Attention (MLA) with FlashInfer-compatible cache. |
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py` at line
1, The file lacks the required NVIDIA copyright header; add the standard NVIDIA
copyright comment block (including the year of latest meaningful modification)
at the very top of
tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py before the
existing module docstring ("""Custom ops for MultiHead Latent Attention (MLA)
with FlashInfer-compatible cache."""), ensuring the header follows the project's
canonical format and includes the correct year.
| @@ -0,0 +1,157 @@ | |||
| """Torch reference implementation for Multi-head Latent Attention (MLA). | |||
Contributor
There was a problem hiding this comment.
Add the required NVIDIA copyright header.
🔧 Proposed fix
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
"""Torch reference implementation for Multi-head Latent Attention (MLA).As per coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of latest meaningful modification.
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py` at line 1, The
file torch_mla.py is missing the required NVIDIA copyright header; prepend the
standard NVIDIA copyright header block (including "Copyright (c) <year>, NVIDIA
CORPORATION. All rights reserved." with the year set to the latest meaningful
modification) at the top of torch_mla.py before the module docstring ("""Torch
reference implementation for Multi-head Latent Attention (MLA).""") so the file
complies with the TensorRT-LLM source file guidelines.
Comment on lines
+346
to
+352
| # Allocate contiguous paged caches (after KVCacheManager so we know page count) | ||
| if self._contiguous_paged_cache_order: | ||
| num_pages = self._kv_cache_manager.blocks_in_primary_pool | ||
| page_size = self._kv_cache_manager.tokens_per_block | ||
| for name, handler in self._contiguous_paged_cache_order.items(): | ||
| self._caches[name] = handler.allocate(self.info, num_pages, page_size) | ||
|
|
Contributor
There was a problem hiding this comment.
Reallocate contiguous paged caches after KVCacheManager resize.
resize_kv_cache_manager() recreates the KVCacheManager and refreshes paged/state views, but contiguous paged caches remain at the old shape. That can cause shape mismatches and invalid accesses after resize.
🛠️ Suggested fix
def resize_kv_cache_manager(self, mem_exclude: int = 0) -> None:
...
max_tokens_final = self._create_kv_cache_manager(max_tokens=max_tokens_optimal)
+
+ # Re-allocate contiguous paged caches to match the new manager size
+ if self._contiguous_paged_cache_order:
+ num_pages = self._kv_cache_manager.blocks_in_primary_pool
+ page_size = self._kv_cache_manager.tokens_per_block
+ for name, handler in self._contiguous_paged_cache_order.items():
+ self._caches[name] = handler.allocate(self.info, num_pages, page_size) def _clear_cache_views(self) -> None:
"""Set paged and state cache views to None before pool release."""
self._kv_cache_config_tuned = None
for name in self._paged_cache_order:
self._caches[name] = None
for name in self._state_resource_order:
self._caches[name] = None
+ for name in self._contiguous_paged_cache_order:
+ self._caches[name] = None🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/shim/interface.py` around lines 346 - 352,
Contiguous paged caches in _contiguous_paged_cache_order are left at the old
shape after resize_kv_cache_manager() recreates _kv_cache_manager, causing shape
mismatches; update the code so that after _kv_cache_manager is resized/recreated
you recompute num_pages = self._kv_cache_manager.blocks_in_primary_pool and
page_size = self._kv_cache_manager.tokens_per_block and then reallocate each
contiguous cache by calling handler.allocate(self.info, num_pages, page_size)
and storing the result back into self._caches[name] (optionally releasing or
overwriting the old cache first) to ensure caches match the new KVCacheManager
shape.
Comment on lines
+73
to
+77
| # if t_name == "compile_model": # or whichever transform is last before generation | ||
| # from tensorrt_llm._torch.auto_deploy.utils.graph_debug_compare import run_comparison | ||
| # run_comparison(mod, cm, self.factory, output_dir="debug_scatter_plots") | ||
| # breakpoint() # Optional: pause to inspect results | ||
|
|
Contributor
There was a problem hiding this comment.
Remove commented-out debug code before merging.
This block contains debugging artifacts (commented-out code with breakpoint()) that should not be committed to the main branch. Additionally, the pipeline failure indicates this section has formatting issues that need to be resolved.
If this debugging capability is valuable, consider:
- Making it conditionally enabled via an environment variable (similar to
AD_DUMP_GRAPHS_DIRpattern used elsewhere in this PR) - Or removing it entirely and re-adding when needed
🗑️ Proposed fix: Remove the commented-out code
mod = transform(mod, cm, self.factory, self.shared_config)
- # if t_name == "compile_model": # or whichever transform is last before generation
- # from tensorrt_llm._torch.auto_deploy.utils.graph_debug_compare import run_comparison
- # run_comparison(mod, cm, self.factory, output_dir="debug_scatter_plots")
- # breakpoint() # Optional: pause to inspect results
-
############################################################################################🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/transform/optimizer.py` around lines 73 - 77,
Remove the commented-out debug block (the lines referencing t_name ==
"compile_model", run_comparison, and breakpoint) from
tensorrt_llm/_torch/auto_deploy/transform/optimizer.py to eliminate debugging
artifacts and fix the formatting error; if you want to preserve the debug
capability instead of deleting it, wrap the logic in a conditional gated by an
environment variable (similar to AD_DUMP_GRAPHS_DIR), e.g., check os.environ and
only call run_comparison(mod, cm, self.factory, output_dir=...) when enabled,
ensuring no leftover commented code or stray breakpoint() calls remain.
Comment on lines
+1
to
+11
| import json | ||
| import os | ||
| import shutil | ||
| from pathlib import Path | ||
| from typing import Any, Dict, Optional | ||
|
|
||
| import dill | ||
| import torch | ||
| import torch.nn as nn | ||
| from safetensors.torch import save_file as safetensors_save | ||
|
|
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if dill and safetensors are declared as dependencies
echo "=== Checking for dill dependency ==="
rg "dill" --glob "setup.py" --glob "pyproject.toml" --glob "requirements*.txt" -A 2 -B 2
echo ""
echo "=== Checking for safetensors dependency ==="
rg "safetensors" --glob "setup.py" --glob "pyproject.toml" --glob "requirements*.txt" -A 2 -B 2
echo ""
echo "=== Listing all dependency configuration files ==="
find . -maxdepth 2 -type f \( -name "setup.py" -o -name "setup.cfg" -o -name "pyproject.toml" -o -name "requirements*.txt" \) | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1883
🏁 Script executed:
#!/bin/bash
echo "=== Content of ./pyproject.toml ==="
cat ./pyproject.toml | head -100
echo ""
echo "=== Content of ./requirements.txt ==="
cat ./requirements.txt
echo ""
echo "=== Content of ./setup.py (first 150 lines) ==="
head -150 ./setup.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 12659
Add dill and safetensors to project dependencies.
The code imports dill (line 7) and safetensors (line 10), but neither package is declared in requirements.txt, setup.py, or pyproject.toml. These external dependencies must be added to the main project dependencies so they are automatically installed with the package, otherwise users will encounter import errors at runtime.
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/utils/logger.py` around lines 1 - 11, The
project imports dill and safetensors in
tensorrt_llm._torch.auto_deploy.utils.logger (the symbols dill and
safetensors.torch.save_file / safetensors_save), but these packages aren't
declared as install-time dependencies; update the project's dependency metadata
(e.g., requirements.txt, setup.py install_requires, or pyproject.toml
[project.dependencies]) to include dill and safetensors (use the appropriate
package names on PyPI) and pin or specify compatible version ranges, then run
dependency install to validate imports succeed.
Comment on lines
+359
to
+392
| class TestGLM4Flash(LlmapiAccuracyTestHarness): | ||
| """Accuracy regression tests for GLM-4.7-Flash.""" | ||
|
|
||
| MODEL_NAME = "zai-org/GLM-4.7-Flash" | ||
| MODEL_PATH = MODEL_NAME # Model is in HF_CACHE | ||
| # Set minimum possible seq len + small buffer, for test speed & memory usage | ||
| MAX_SEQ_LEN = max(MMLU.MAX_INPUT_LEN + MMLU.MAX_OUTPUT_LEN, | ||
| GSM8K.MAX_INPUT_LEN + GSM8K.MAX_OUTPUT_LEN) | ||
| MAX_NUM_TOKENS = MAX_SEQ_LEN | ||
|
|
||
| def get_default_kwargs(self, enable_chunked_prefill=False): | ||
| config = { | ||
| "skip_tokenizer_init": False, | ||
| "trust_remote_code": True, | ||
| "compile_backend": "torch-cudagraph", | ||
| "max_batch_size": 128, | ||
| "max_seq_len": self.MAX_SEQ_LEN, | ||
| "max_num_tokens": self.MAX_NUM_TOKENS, | ||
| "skip_loading_weights": False, | ||
| "disable_overlap_scheduler": False, | ||
| "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128], | ||
| "kv_cache_config": { | ||
| "enable_block_reuse": False, | ||
| "free_gpu_memory_fraction": 0.88 | ||
| }, | ||
| "model_kwargs": { | ||
| "torch_dtype": "bfloat16" | ||
| }, | ||
| } | ||
| if enable_chunked_prefill: | ||
| config["enable_chunked_prefill"] = True | ||
| config[ | ||
| "max_num_tokens"] = 512 # NOTE: must be > max(tokens_per_block, max_batch_size) | ||
| return config |
Contributor
There was a problem hiding this comment.
Inconsistent disable_overlap_scheduler setting.
The test sets disable_overlap_scheduler: False (line 378), but the YAML config at examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml sets it to true. Consider aligning these values for consistency, or document if the difference is intentional for testing purposes.
🤖 Prompt for AI Agents
In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py` around lines 359
- 392, The disable_overlap_scheduler flag in TestGLM4Flash.get_default_kwargs is
set to False but differs from the model registry's configuration; update
get_default_kwargs to match the registry flag (set "disable_overlap_scheduler":
True) to keep behavior consistent, or explicitly document the intentional
mismatch by adding a clear comment in the TestGLM4Flash class and in
get_default_kwargs noting why tests require False while the deployed model uses
True; refer to the TestGLM4Flash class and get_default_kwargs to locate the
change and to the "disable_overlap_scheduler" key to update or document.
Comment on lines
+485
to
+523
| # Debug: print state dict keys and shapes | ||
| print("\n=== HF MoE state_dict keys and shapes ===") | ||
| for k, v in hf_state_dict.items(): | ||
| print(f" {k}: {v.shape}") | ||
|
|
||
| custom_state_dict = _convert_hf_moe_state_dict_to_custom(hf_state_dict, config.n_routed_experts) | ||
|
|
||
| print("\n=== Converted custom state_dict keys and shapes ===") | ||
| for k, v in custom_state_dict.items(): | ||
| print(f" {k}: {v.shape}") | ||
|
|
||
| print("\n=== Expected custom MoE state_dict keys ===") | ||
| for k, v in custom_moe.state_dict().items(): | ||
| print(f" {k}: {v.shape}") | ||
|
|
||
| custom_moe.load_state_dict(custom_state_dict) | ||
| custom_moe.eval() | ||
|
|
||
| # Sanity check: verify expert weights match after conversion | ||
| # HF has stacked weights: experts.gate_up_proj [n_experts, 2*intermediate, hidden] | ||
| # Our model has per-expert: experts.{i}.gate_proj.weight, experts.{i}.up_proj.weight | ||
| hf_gate_up = hf_moe.experts.gate_up_proj # [n_experts, 2*intermediate, hidden] | ||
| hf_down = hf_moe.experts.down_proj # [n_experts, hidden, intermediate] | ||
| intermediate_size = config.moe_intermediate_size | ||
|
|
||
| print(f"\n=== Debug: intermediate_size = {intermediate_size} ===") | ||
| print(f"hf_gate_up shape: {hf_gate_up.shape}") | ||
| print(f"hf_gate_up[0, :2, :2]: {hf_gate_up[0, :2, :2]}") | ||
|
|
||
| # Get the converted state dict values for comparison | ||
| converted_gate_0 = custom_state_dict["experts.0.gate_proj.weight"] | ||
| print(f"converted_gate_0 shape: {converted_gate_0.shape}") | ||
| print(f"converted_gate_0[:2, :2]: {converted_gate_0[:2, :2]}") | ||
|
|
||
| # After load_state_dict | ||
| loaded_gate_0 = custom_moe.experts[0].gate_proj.weight | ||
| print(f"loaded_gate_0 shape: {loaded_gate_0.shape}") | ||
| print(f"loaded_gate_0[:2, :2]: {loaded_gate_0[:2, :2]}") | ||
|
|
Contributor
There was a problem hiding this comment.
Remove debug print statements from the numerical equivalence test.
These prints will spam CI logs and slow test runs. Prefer removing them or gating behind an explicit debug flag.
🧹 Proposed cleanup
- # Debug: print state dict keys and shapes
- print("\n=== HF MoE state_dict keys and shapes ===")
- for k, v in hf_state_dict.items():
- print(f" {k}: {v.shape}")
-
custom_state_dict = _convert_hf_moe_state_dict_to_custom(hf_state_dict, config.n_routed_experts)
-
- print("\n=== Converted custom state_dict keys and shapes ===")
- for k, v in custom_state_dict.items():
- print(f" {k}: {v.shape}")
-
- print("\n=== Expected custom MoE state_dict keys ===")
- for k, v in custom_moe.state_dict().items():
- print(f" {k}: {v.shape}")
-
custom_moe.load_state_dict(custom_state_dict)
custom_moe.eval()
@@
- print(f"\n=== Debug: intermediate_size = {intermediate_size} ===")
- print(f"hf_gate_up shape: {hf_gate_up.shape}")
- print(f"hf_gate_up[0, :2, :2]: {hf_gate_up[0, :2, :2]}")
-
# Get the converted state dict values for comparison
converted_gate_0 = custom_state_dict["experts.0.gate_proj.weight"]
- print(f"converted_gate_0 shape: {converted_gate_0.shape}")
- print(f"converted_gate_0[:2, :2]: {converted_gate_0[:2, :2]}")
-
# After load_state_dict
loaded_gate_0 = custom_moe.experts[0].gate_proj.weight
- print(f"loaded_gate_0 shape: {loaded_gate_0.shape}")
- print(f"loaded_gate_0[:2, :2]: {loaded_gate_0[:2, :2]}")🤖 Prompt for AI Agents
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_glm4_moe_lite_modeling.py`
around lines 485 - 523, Remove the debug print statements in the numerical
equivalence test: delete or gate all prints that reference hf_state_dict,
custom_state_dict, custom_moe, hf_gate_up, converted_gate_0, and loaded_gate_0
(the blocks printing shapes, slices, and the "=== ..." headers). Keep the actual
assertions and the call to _convert_hf_moe_state_dict_to_custom and
custom_moe.load_state_dict as-is; if logging is desired for debugging, wrap the
prints behind an explicit debug flag (e.g., a module-level DEBUG or an env var)
so CI runs remain quiet.
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
bmarimuthu-nv
force-pushed
the
bala/glm4.7-flash-standalone
branch
from
February 5, 2026 09:15
a92d586 to
affa2ed
Compare
1 task
Contributor
|
see #11324 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes
New Features
Documentation
Bug Fixes
fixes:
#11032
#11113
Description
Test Coverage
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)
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
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.