[Draft] AutoDeploy GLM4.7 flash bundle - #11356
Conversation
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Co-authored-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> Co-authored-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
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: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
📝 WalkthroughWalkthroughThis pull request adds support for GLM-4.7-Flash model deployment, introduces new DeepSeekV3 and GLM4 MoE Lite model implementations, refactors the MLA attention system with FlashInfer and Torch backends, adds interleaved RoPE optimization, and expands test coverage across multiple components. Changes span model implementations, custom operations, configurations, and comprehensive test suites. Changes
Sequence Diagram(s)sequenceDiagram
actor Client
participant Torch as Torch MLA
participant FlashInfer as FlashInfer Planner
participant CudaGraph as CUDA Graph
participant Kernel as GPU Kernel
rect rgba(100, 200, 100, 0.5)
Note over Client,Kernel: Prefill Phase
Client->>Torch: torch_mla(q, k_expanded)
Torch->>Kernel: Compute attention
Kernel-->>Torch: Output
Torch-->>FlashInfer: Cache data
end
rect rgba(100, 100, 200, 0.5)
Note over Client,Kernel: Generate Phase (Absorption)
Client->>FlashInfer: Plan decode wrapper
FlashInfer->>CudaGraph: Warm-up CUDA graph
CudaGraph->>Kernel: Build kernel launch sequence
Client->>CudaGraph: torch_backend_mla_with_cache
CudaGraph->>Kernel: Replayed graph
Kernel-->>CudaGraph: Output
CudaGraph-->>Client: Result
end
rect rgba(200, 100, 100, 0.5)
Note over Client,Kernel: Decode with FlashInfer Cache
Client->>FlashInfer: flashinfer_mla_with_cache
FlashInfer->>Kernel: Paged attention op
Kernel-->>FlashInfer: Output projection
FlashInfer-->>Client: Attention output
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ 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}
There was a problem hiding this comment.
Actionable comments posted: 12
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py (1)
31-43:⚠️ Potential issue | 🟡 MinorTest has no assertions — it will always pass.
The test only prints cache values but never asserts correctness. After calling
_update_kv_cache, verify the expected state ofk_cacheandv_cache. For instance, slots 0–2 ofk_cache[0]should be filled with ones andk_cache[1][0]should hold the last token.🧪 Proposed fix: add assertions
_update_kv_cache( k.view(batch_size * seq_length, n_heads, K_D_HEAD), v.view(batch_size * seq_length, n_heads, V_D_HEAD), k_cache, v_cache, torch.tensor([3, 1]).long(), torch.tensor([0, 0]), slot_idx=torch.tensor([0, 1]), seq_start=torch.tensor([0, 3]).long(), ) - print("k_cache: " + str(k_cache)) - print("v_cache: " + str(v_cache)) + # Sequence 0: 3 tokens starting at seq_start=0, written to slot 0, positions 0:3 + assert torch.equal(k_cache[0, :3, :, :], torch.ones(3, n_heads, K_D_HEAD)) + assert torch.equal(k_cache[0, 3, :, :], torch.zeros(n_heads, K_D_HEAD)) + # Sequence 1: 1 token starting at seq_start=3, written to slot 1, position 0 + assert torch.equal(k_cache[1, 0, :, :], torch.ones(n_heads, K_D_HEAD)) + assert torch.equal(k_cache[1, 1:, :, :], torch.zeros(MAX_SEQ_LEN - 1, n_heads, K_D_HEAD)) + # Same checks for v_cache + assert torch.equal(v_cache[0, :3, :, :], torch.ones(3, n_heads, V_D_HEAD)) + assert torch.equal(v_cache[1, 0, :, :], torch.ones(n_heads, V_D_HEAD))
🤖 Fix all issues with AI agents
In `@examples/auto_deploy/model_registry/models.yaml`:
- Around line 224-225: The registry entry for name "zai-org/GLM-4.7-Flash" is
missing shared runtime configuration; update its yaml_extra list to include
"dashboard_default.yaml" (to supply runtime: trtllm, attn_backend: flashinfer,
model_factory: AutoModelForCausalLM, skip_loading_weights: false) and add an
appropriate "world_size_N.yaml" (e.g., "world_size_1.yaml" or the correct
cluster world size file used by other models) so the deployment will inherit the
required runtime and world-size-specific settings.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py`:
- Around line 11-13: Add the required NVIDIA copyright header at the top of this
__init__.py (before any imports) using the exact SPDX and Apache-2.0 block
provided in the review, preserving the existing imports (FlashInferMLAAttention,
flashinfer_mla_with_cache, TorchBackendMLAAttention,
torch_backend_mla_with_cache, torch_mla); ensure the header appears as the very
first lines of the file and is not commented out incorrectly so tooling
recognizes the SPDX tags.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py`:
- Around line 1-27: Add the required NVIDIA copyright header to the top of this
module file (tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py)
above the existing module docstring; ensure the header matches the project's
standard copyright template and includes year and owner, then keep the existing
docstring and imports (FlashInferMLAAttention, flashinfer_mla_with_cache, etc.)
unchanged below it.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py`:
- Around line 1-23: Add the required NVIDIA copyright header at the top of this
module file (before the module docstring) per project guidelines; update the
file containing torch_cached_mla_with_cache and TorchBackendMLAAttention to
prepend the standard NVIDIA comment block (including copyright year(s) and
"NVIDIA CORPORATION" wording, license reference if required by repo policy) so
every .py source file includes the header.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py`:
- Around line 1-12: This file is missing the required NVIDIA copyright header —
add the standard NVIDIA copyright header (including the year of latest
meaningful modification) at the very top of
tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py before the module
docstring; ensure the header follows the project's header format used across
other source files and includes the appropriate year and ownership text so the
module (torch_mla.py) complies with the repository coding guidelines.
- Around line 109-115: The code currently only applies the causal mask when s_q
== s_k, silently skipping it otherwise; change the logic in torch_mla.py so that
when is_causal is True you always build a rectangular causal_mask =
torch.triu(torch.ones(s_q, s_k, device=q_nope.device, dtype=torch.bool),
diagonal=1) (i.e., remove the s_q == s_k guard) and apply it to attn_scores via
masked_fill_(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"));
alternatively, if you prefer to forbid mismatched lengths, explicitly raise a
ValueError when is_causal and s_q != s_k instead—use the variables is_causal,
s_q, s_k, q_nope, and attn_scores to locate and implement the change.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope_kernel.py`:
- Around line 146-161: The kernel computes q/k output offsets assuming
contiguous [B,S,H,D] layout (see q_out_stride_b, q_out_stride_s, q_out_stride_h,
q_out_offsets_first and the tl.store calls to q_out_ptr), but outputs are
allocated with torch.empty_like which may preserve a non-contiguous memory
format; fix by ensuring inputs and outputs are contiguous before launching the
Triton kernel: make q and k contiguous (or allocate outputs explicitly with
memory_format=torch.contiguous_format) in the wrapper that prepares
q_out_ptr/k_out_ptr so the calculated strides/offsets match the actual memory
layout used by tl.store.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py`:
- Line 1: Update the copyright header in the modeling_deepseek.py source file
from 2025 to 2026; locate the top-of-file copyright comment in
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py and change
the year to 2026 so the header reflects the latest meaningful modification year.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py`:
- Line 1: Update the file header comment that currently reads "Copyright (c)
2025, NVIDIA CORPORATION. All rights reserved." to reflect the latest
modification year 2026; replace "2025" with "2026" in the top-of-file copyright
string in
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py so the
header becomes "Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved."
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py`:
- Around line 2134-2145: The test
test_flashinfer_mla_init_decode_wrapper_with_buffers declares a parameterized
dtype but never uses it; either remove the `@pytest.mark.parametrize`("dtype",
[torch.bfloat16]) line or apply the dtype when constructing buffer tensors in
the test so the parameter is exercised (inspect the tensor creation sites in
test_flashinfer_mla_init_decode_wrapper_with_buffers and update them to use the
parameter or delete the unused parametrize), referencing the test function name
and the _GlobalFlashInferMLAPlanner setup to locate the code.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_torch_mla_op.py`:
- Around line 33-34: Add Optional to the typing imports and change the two
parameters that default to None to explicit Optional types: replace the implicit
"scale: float = None" with "scale: Optional[float]" and "kv_lora_rank: int =
None" with "kv_lora_rank: Optional[int]" in the function or signature where
these parameters are declared (look for the parameters named scale and
kv_lora_rank in the test_torch_mla_op.py signature) so the type hints comply
with PEP 484 and Ruff RUF013.
🧹 Nitpick comments (24)
tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py (1)
397-414: Consider increasinghidden_sizeto be more proportionate with the new MLA constraints.With
kv_lora_rank: 512andqk_rope_head_dim: 64, these dimensions dwarf thehidden_size: 32(ratios of 16× and 2× respectively). In the real DeepSeek-V3 model,kv_lora_rankis ~7% ofhidden_size. This imbalance could make projections (e.g., the down-projection fromkv_lora_rank→hidden_size) behave unrealistically and may inflate memory/compute in tests relative to the other "small" configs.If flashinfer_mla mandates these minimums, consider bumping
hidden_size(and related dims likeintermediate_size,q_lora_rank) upward so the overall config stays internally consistent — even if it's still much smaller than the real model. For example,hidden_size: 512would restore a more realistic ratio while remaining small for testing.That said, the comments clearly document the constraint, and if the current config already passes tests at acceptable speed, this is fine to defer.
tensorrt_llm/_torch/auto_deploy/config/default.yaml (1)
163-165: Nit: Comment separator length inconsistency.Lines 163 and 165 use 27
#characters while all other section separators in this file use 28. This is cosmetic only.- ########################################################################################### + ############################################################################################ # SWITCH TO CACHED+FLATTENED ATTENTION + INITIALIZE CACHES - ########################################################################################### + ############################################################################################tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py (1)
147-157: Both branches of the layout check intorch_mla_fakeproduce identical code.The
if/elseblocks at lines 148-156 are identical. Sinceq_nope.shape[0:3]already encodes the correct layout-dependent dimensions, the branch can be collapsed.Suggested simplification
# Output shape depends on layout - if layout == "bsnd": - # Input: [B, S, N, D], Output: [B, S, N, v_head_dim] - return q_nope.new_empty( - q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim - ).contiguous() - else: - # Input: [B, N, S, D], Output: [B, N, S, v_head_dim] - return q_nope.new_empty( - q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim - ).contiguous() + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous()tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py (4)
47-64: Custom__hash__via string concatenation works but is fragile.The hash implementation at line 64 (
hash("_".join([str(getattr(self, f.name)) for f in fields(self)]))) joins string representations with_. If any field's string representation itself contains_, two different parameter sets could theoretically collide (e.g.,dtype=torch.float16withsm_scale=Nonevs different combinations). This is unlikely with the current field types but atuple-based hash would be more robust:def __hash__(self): return hash(tuple(getattr(self, f.name) for f in fields(self)))This applies to both
MLADecodePlanParamsandMLAPrefillPlanParams.Also applies to: 67-82
141-162:self.__init__()as a reset mechanism is unconventional.Line 149 calls
self.__init__()to reinitialize all state. This works but is unusual and can confuse readers. A dedicated_clear_state()helper would be more explicit.
570-631: Duplicatedkv_b_proj_weightdecomposition in chunked prefill and decode paths.Lines 579-585 (chunked prefill) and lines 705-711 (decode) perform identical reshaping and splitting of
kv_b_proj_weightintow_knandw_v. Extract this into a small helper to reduce duplication.Suggested helper
def _decompose_kv_b_proj(kv_b_proj_weight, num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank): """Split kv_b_proj_weight into W_kn and W_v components.""" reshaped = kv_b_proj_weight.view(num_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank) w_kn = reshaped[:, :qk_nope_head_dim, :] w_v = reshaped[:, qk_nope_head_dim:, :] return w_kn, w_vAlso applies to: 692-717
929-933: Enhance error messages with actual values and FlashInfer documentation link.The checks
kv_lora_rank != 512andqk_rope_head_dim != 64enforce confirmed FlashInfer MLA library requirements. Improve the error messages to show the actual received values and link to FlashInfer documentation:Suggested error messages
raise ValueError( f"FlashInfer MLA requires kv_lora_rank=512, got {kv_lora_rank}. " "See https://docs.flashinfer.ai/api/mla.html" )raise ValueError( f"FlashInfer MLA requires qk_rope_head_dim=64, got {qk_rope_head_dim}. " "See https://docs.flashinfer.ai/api/mla.html" )tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py (1)
340-345: Parameter shadowing:seq_len,input_pos, andslot_idxare reassigned inside the function.Lines 342-344 reassign the function parameters
seq_len,input_pos, andslot_idxto sliced versions, shadowing the original values. Per coding guidelines, avoid shadowing variables declared in an outer scope. Consider using distinct names likeseq_len_trimmed, or applying slicing inline where these are used.Suggested fix
- num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_prefill, _num_prefill_tokens, num_decode = batch_info_host.tolist() num_seq = num_prefill + num_decode - seq_len = seq_len[:num_seq] - input_pos = input_pos[:num_seq] - slot_idx = slot_idx[:num_seq] - seq_start = cu_seqlen[:num_seq] + seq_len_active = seq_len[:num_seq] + input_pos_active = input_pos[:num_seq] + slot_idx_active = slot_idx[:num_seq] + seq_start = cu_seqlen[:num_seq]Then update references to
seq_len,input_pos,slot_idxin the rest of the function accordingly.tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py (2)
364-365: Prefix unused unpacked variables with_.
bszandseq_lenare unpacked but unused. Prefix them with underscores for clarity.Proposed fix
- bsz, seq_len, hidden_dim = hidden_states.shape + _bsz, _seq_len, hidden_dim = hidden_states.shape
684-737: Rope-scaling config checks are repeated in three places.The
config.rope_scaling is not None and isinstance(..., dict) and "factor" in ...pattern appears in_init_rope(lines 690–694 and 704–708) and inGlm4MoeLiteAttention.__init__(lines 504–508). Consider extracting a small helper (e.g.,_has_yarn_scaling(config)) to consolidate this.tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_glm4_moe_lite_modeling.py (4)
485-522: Remove debug print statements.Lines 486–498 and 510–522 contain development-time
print()calls that add noise to CI output. Either remove them or, if needed for debugging failures, convert tologging.debug().
562-563: Consider documenting why 5% tolerance is used.
rtol=0.05, atol=0.05is very loose for a numerical equivalence test. If this is justified by fused-vs-unfused routing causing different expert selections at decision boundaries, a brief comment would help future maintainers understand this isn't masking real issues.
302-311:_get_hf_attention_classappears unused.This helper is defined but not called by any test function in this file. Remove it to avoid dead code, or add the attention equivalence test it was intended for.
406-451: Config parameter duplication between_create_small_configand_create_hf_config.Both functions define identical parameter values separately. If one is updated without the other, numerical equivalence tests will silently test mismatched configs. Consider extracting the shared parameters into a module-level dict.
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py (3)
529-534:config_classis not set onDeepSeekV3PreTrainedModel.Unlike
Glm4MoeLitePreTrainedModelwhich setsconfig_class = Glm4MoeLiteConfig, the DeepSeek variant omits it. This meansfrom_pretrained()/save_pretrained()won't resolve the config class correctly if used directly. If this is intentional (auto_deploy-only usage), a brief comment would clarify the design choice.
85-170: Substantial code duplication withmodeling_glm4_moe_lite.py.
DeepSeekV3YarnRotaryEmbedding(including all YaRN static methods),DeepSeekV3MLP,DeepSeekV3MoE, output dataclasses, and the model scaffolding are structurally near-identical to their GLM4 counterparts. Consider extracting shared building blocks (at minimum the YaRN utilities, MLP, and output dataclasses) into a common module to reduce maintenance burden.This is understandable for a draft bundle PR, but worth addressing before the individual upstream PRs merge.
355-362: Less defensiverope_scalingaccess compared to GLM4 counterpart.Line 357 accesses
config.rope_scaling["factor"]and line 372 accesses["type"]without checking key existence. The GLM4 file guards withisinstance(config.rope_scaling, dict) and "factor" in config.rope_scaling. Low risk for standard HF configs, but inconsistent and could produce confusingKeyErroron malformed configs.tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py (1)
287-287: Pre-existing:mutates_args=()may be inaccurate.The custom op mutates
k_cacheandv_cachein-place (via_update_kv_cacheand_torch_generate_mha), butmutates_args=()declares no mutations. This isn't introduced by this change, but worth a follow-up since incorrect mutation declarations can cause subtle issues withtorch.compileand FX tracing.examples/auto_deploy/cookbooks/glm_4.7_flash_trtllm_cookbook.ipynb (1)
74-95: Consider clearing execution outputs before committing.The notebook includes hardcoded execution outputs from a specific 8×H100 environment (Python version, GPU names, counts). These will differ for every user and add noise to diffs. Clearing outputs keeps the notebook portable and reduces repository churn.
tensorrt_llm/_torch/auto_deploy/transform/library/rope.py (1)
617-617: Unusedrest— prefix with underscore.Ruff correctly flags this. The
*reston line 617 captures unused trailing args.🔧 Fix
- q_node, k_node, cos_node, sin_node, *rest = node.args + q_node, k_node, cos_node, sin_node, *_rest = node.argstests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_torch_mla_op.py (2)
124-128: Redundant conditional — both branches are identical.The
if is_generate/elseblock on Lines 124–127 produces the exact samequery_fullin both branches. This should be simplified.Proposed fix
- # Construct full query and key - if is_generate: - query_full = np.concatenate([q_nope_seq, q_pe_seq], axis=-1) - else: - query_full = np.concatenate([q_nope_seq, q_pe_seq], axis=-1) + # Construct full query and key + query_full = np.concatenate([q_nope_seq, q_pe_seq], axis=-1)
655-665: Subtle ordering: numpy reference receives cache after the torch op already mutated it.
self._run_cached_mla(data)on Line 652 mutatesdata["mla_cache"]in-place. The numpy reference on Line 655 then receives a copy of the already-updated cache. This works because the numpy reference overwrites the same positions with the same data, but the coupling is fragile — if the torch op or reference ever diverge in how they write to the cache, this test silently passes with inconsistent state.Consider cloning the cache before running the torch op and passing the clone to the numpy reference, so both start from the same pre-mutation state.
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py (2)
636-892: Large parameter space may cause GPU OOM for some configurations.With
batch_size=256andprefill_seq_length=1024, the paged caches alone (max_num_pages≈ 8448 pages × 64 tokens × 512 dims × 2 bytes) will consume ~550 MB forckv_cachealone. Combined with all other tensors and the fact that this is a parameterized test, this could be tight on CI GPUs.Consider whether
batch_size=256withprefill_seq_length=1024is necessary, or whether a smaller combination would still provide sufficient coverage.
1354-1357: Unguarded assumption: all sequences have the same number of chunks.Line 1356 uses
len(chunks_per_seq[0])and iteratesrange(num_chunks)for all sequences. If a future parametrize config has sequences with differing chunk counts, Line 1396 (chunks_per_seq[seq_idx][chunk_idx]) will raise anIndexError. Consider adding a guard assertion.Proposed fix
chunks_per_seq = chunk_config["chunks_per_seq"] batch_size = len(chunks_per_seq) num_chunks = len(chunks_per_seq[0]) # Assume all sequences have same number of chunks + assert all(len(c) == num_chunks for c in chunks_per_seq), ( + "All sequences must have the same number of chunks" + )
| - name: zai-org/GLM-4.7-Flash | ||
| yaml_extra: ['glm-4.7-flash.yaml'] |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, let's see the full context of models.yaml around lines 224-225
echo "=== Context from models.yaml (lines 210-230) ==="
sed -n '210,230p' examples/auto_deploy/model_registry/models.yaml
echo ""
echo "=== Check a few other model entries for pattern ==="
# Look at several model entries to confirm the pattern
sed -n '1,100p' examples/auto_deploy/model_registry/models.yaml | grep -A 1 "- name:"
echo ""
echo "=== Look for dashboard_default.yaml ==="
find examples/auto_deploy/model_registry -name 'dashboard_default.yaml' -type f
echo ""
echo "=== Look for world_size configs ==="
find examples/auto_deploy/model_registry -name 'world_size*.yaml' -type f | head -5
echo ""
echo "=== Check if glm-4.7-flash.yaml exists ==="
find examples/auto_deploy/model_registry -name 'glm-4.7-flash.yaml' -type fRepository: NVIDIA/TensorRT-LLM
Length of output: 1772
🏁 Script executed:
#!/bin/bash
echo "=== dashboard_default.yaml content ==="
cat examples/auto_deploy/model_registry/configs/dashboard_default.yaml
echo ""
echo "=== glm-4.7-flash.yaml content ==="
cat examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml
echo ""
echo "=== Check world_size_8.yaml for comparison ==="
cat examples/auto_deploy/model_registry/configs/world_size_8.yaml
echo ""
echo "=== Count all models with this pattern to confirm ==="
grep -c "yaml_extra:" examples/auto_deploy/model_registry/models.yaml
echo "Models with dashboard_default.yaml:"
grep "dashboard_default.yaml" examples/auto_deploy/model_registry/models.yaml | wc -l
echo "Models without dashboard_default.yaml:"
grep -v "dashboard_default.yaml" examples/auto_deploy/model_registry/models.yaml | grep "yaml_extra:" | wc -lRepository: NVIDIA/TensorRT-LLM
Length of output: 808
GLM-4.7-Flash entry is missing critical runtime configuration from dashboard_default.yaml.
This entry omits 'dashboard_default.yaml' and 'world_size_*.yaml', which are included in 103 out of 104 models. The glm-4.7-flash.yaml file does not provide essential parameters that dashboard_default.yaml defines: runtime (trtllm), attn_backend (flashinfer), model_factory (AutoModelForCausalLM), and skip_loading_weights (false). Without these, the deployment registry will lack critical runtime configuration. Add 'dashboard_default.yaml' and an appropriate 'world_size_N.yaml' to the yaml_extra list.
🤖 Prompt for AI Agents
In `@examples/auto_deploy/model_registry/models.yaml` around lines 224 - 225, The
registry entry for name "zai-org/GLM-4.7-Flash" is missing shared runtime
configuration; update its yaml_extra list to include "dashboard_default.yaml"
(to supply runtime: trtllm, attn_backend: flashinfer, model_factory:
AutoModelForCausalLM, skip_loading_weights: false) and add an appropriate
"world_size_N.yaml" (e.g., "world_size_1.yaml" or the correct cluster world size
file used by other models) so the deployment will inherit the required runtime
and world-size-specific settings.
| from .flashinfer_mla import FlashInferMLAAttention, flashinfer_mla_with_cache | ||
| from .torch_backend_mla import TorchBackendMLAAttention, torch_backend_mla_with_cache | ||
| from .torch_mla import torch_mla |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, examine the flashinfer_mla.py file to confirm the unconditional import
head -30 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 1339
🏁 Script executed:
# Check the mla __init__.py file itself
cat -n tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1017
🏁 Script executed:
# Check for optional dependency patterns in the codebase
rg "try:" -A 2 -B 0 --type=py | grep -A 2 "ImportError" | head -30Repository: NVIDIA/TensorRT-LLM
Length of output: 1150
🏁 Script executed:
# Check setup.py or pyproject.toml for flashinfer dependency info
fd -e "setup.py" -e "pyproject.toml" -e "requirements.txt" | xargs grep -l "flashinfer" 2>/dev/nullRepository: NVIDIA/TensorRT-LLM
Length of output: 47
🏁 Script executed:
# Look for how flashinfer is handled in other files
rg "flashinfer" --type=py -l | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1520
🏁 Script executed:
# Search for dependency config files and flashinfer mentions
find . -maxdepth 2 -name "pyproject.toml" -o -name "setup.py" -o -name "setup.cfg" | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 125
🏁 Script executed:
# Check if flashinfer is mentioned in dependency files
rg "flashinfer" -g "pyproject.toml" -g "setup.py" -g "setup.cfg" -g "requirements*.txt"Repository: NVIDIA/TensorRT-LLM
Length of output: 680
🏁 Script executed:
# Check how other similar backend modules handle optional imports
find tensorrt_llm/_torch/auto_deploy/custom_ops -name "__init__.py" | head -5 | xargs cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 5915
🏁 Script executed:
# Check if the parent auto_deploy module or other backends conditionally import things
rg "try:" -A 5 --type=py tensorrt_llm/_torch/auto_deploy/custom_ops/ | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 2396
🏁 Script executed:
# Check if flashinfer is documented as optional anywhere
rg -i "optional.*flashinfer|flashinfer.*optional" --type=py --type=md | head -10Repository: NVIDIA/TensorRT-LLM
Length of output: 1080
🏁 Script executed:
# Check main pyproject.toml for flashinfer dependency definition
cat pyproject.toml | grep -A 20 "dependencies"Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check if flashinfer is in optional dependencies
cat pyproject.toml | grep -A 30 "optional-dependencies"Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check if the mla __init__.py has a copyright header (from guidelines)
head -20 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 868
🏁 Script executed:
# Check torch_mla.py to see if it also unconditionally imports (to understand the full import chain)
head -40 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 2256
🏁 Script executed:
# Check torch_backend_mla.py
head -40 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 1575
🏁 Script executed:
# Get the actual pyproject.toml content to check dependencies
cat pyproject.toml | head -150Repository: NVIDIA/TensorRT-LLM
Length of output: 7741
🏁 Script executed:
# Check if there are conditional imports for flashinfer elsewhere
rg "try.*import.*flashinfer" --type=py -A 3Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look at the parent custom_ops __init__.py to see if it does dynamic imports
cat -n tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 1510
🏁 Script executed:
# Check dependencies section in main pyproject.toml - search more thoroughly
rg "dependencies|flashinfer" pyproject.toml | head -40Repository: NVIDIA/TensorRT-LLM
Length of output: 363
🏁 Script executed:
# Check if there's an optional-dependencies section that includes flashinfer
rg -A 5 "optional-dependencies" pyproject.toml | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# See how the module is typically used - check if there are conditional imports of mla elsewhere
rg "from.*mla|import.*mla" --type=py | grep -v test | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1072
🏁 Script executed:
# Check if the parent auto_deploy init guards any imports
cat tensorrt_llm/_torch/auto_deploy/__init__.py | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 417
🏁 Script executed:
# Check if flashinfer is listed as required or optional in requirements.txt or setup.py
cat requirements.txt | grep -i flashinferRepository: NVIDIA/TensorRT-LLM
Length of output: 87
🏁 Script executed:
# Check if there's a distinction between required and optional flashinfer usage
cat setup.py | grep -A 10 -B 5 flashinferRepository: NVIDIA/TensorRT-LLM
Length of output: 631
🏁 Script executed:
# Verify what the parent custom_ops init does - it walks and imports everything
# This means if mla/__init__.py imports flashinfer, it's unconditional at module level
cat -n tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1510
🏁 Script executed:
# Check if flashinfer-python is in setup.py as a required dependency
cat setup.py | head -200 | grep -A 50 "install_requires"Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check the copyright header requirement - verify what other __init__.py files in custom_ops have
find tensorrt_llm/_torch/auto_deploy/custom_ops -name "__init__.py" -exec head -5 {} + | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 2576
🏁 Script executed:
# Verify the actual state of the mla/__init__.py file to see if it has copyright
head -20 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 868
🏁 Script executed:
# Check the year in other __init__.py copyright headers to see current pattern
grep "SPDX-FileCopyrightText" tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 166
Add NVIDIA copyright header to file.
The file is missing the required NVIDIA copyright header. Since this file was substantially rewritten, add the header at the beginning:
# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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.The unconditional import of flashinfer_mla is correct since flashinfer-python is a required dependency (specified in requirements.txt).
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py` around lines 11 -
13, Add the required NVIDIA copyright header at the top of this __init__.py
(before any imports) using the exact SPDX and Apache-2.0 block provided in the
review, preserving the existing imports (FlashInferMLAAttention,
flashinfer_mla_with_cache, TorchBackendMLAAttention,
torch_backend_mla_with_cache, torch_mla); ensure the header appears as the very
first lines of the file and is not commented out incorrectly so tooling
recognizes the SPDX tags.
| """FlashInfer-based MLA (Multi-head Latent Attention) backend with paged caching. | ||
|
|
||
| This module provides: | ||
| - FlashInferMLAAttention: attention descriptor using FlashInfer MLA kernels | ||
| - flashinfer_mla_with_cache: cached backend op with paged KV cache | ||
|
|
||
| FlashInfer MLA uses: | ||
| - Regular prefill (input_pos == 0): BatchPrefillWithRaggedKVCacheWrapper with expanded K, V | ||
| - Chunked prefill (input_pos > 0): BatchMLAPagedAttentionWrapper with matrix absorption | ||
| - Decode: BatchMLAPagedAttentionWrapper with paged compressed KV cache | ||
|
|
||
| FlashInfer MLA Cache Layout (two separate caches): | ||
| ckv_cache: [num_pages, page_size, kv_lora_rank] | ||
| kpe_cache: [num_pages, page_size, qk_rope_head_dim] | ||
| - No num_heads dimension (MLA-specific optimization) | ||
|
|
||
| Reference: https://docs.flashinfer.ai/api/mla.html | ||
| """ | ||
|
|
||
| import math | ||
| from dataclasses import dataclass, fields | ||
| from math import prod | ||
| from typing import Dict, List, Literal, Optional, Tuple | ||
|
|
||
| import flashinfer | ||
| import torch | ||
| from torch._ops import OpOverloadPacket |
There was a problem hiding this comment.
Missing NVIDIA copyright header.
As per coding guidelines, all .py source files should contain an NVIDIA copyright header.
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` around
lines 1 - 27, Add the required NVIDIA copyright header to the top of this module
file (tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py) above
the existing module docstring; ensure the header matches the project's standard
copyright template and includes year and owner, then keep the existing docstring
and imports (FlashInferMLAAttention, flashinfer_mla_with_cache, etc.) unchanged
below it.
| # check if we are in cuda graph capture and just return the pre-cached decode wrapper | ||
| if torch.cuda.is_current_stream_capturing() or cuda_graph_state.in_warm_up(): | ||
| wrapper = self.cached_cuda_graph_decode_wrappers[plan_params] | ||
| return wrapper |
There was a problem hiding this comment.
Potential KeyError if CUDA graph capture runs without matching warm-up.
At line 276, self.cached_cuda_graph_decode_wrappers[plan_params] is accessed during CUDA graph capture or warm-up. If the warm-up phase didn't cache a wrapper for the current plan_params (e.g., unexpected batch size), this raises KeyError with no actionable error message.
Suggested defensive check
# check if we are in cuda graph capture and just return the pre-cached decode wrapper
if torch.cuda.is_current_stream_capturing() or cuda_graph_state.in_warm_up():
- wrapper = self.cached_cuda_graph_decode_wrappers[plan_params]
+ if plan_params not in self.cached_cuda_graph_decode_wrappers:
+ raise RuntimeError(
+ f"No cached CUDA graph wrapper for MLA decode with params: {plan_params}. "
+ "Ensure warm-up covers all expected batch sizes."
+ )
+ wrapper = self.cached_cuda_graph_decode_wrappers[plan_params]
return wrapper| """Custom ops for MultiHead Latent Attention (MLA) with FlashInfer-compatible cache. | ||
|
|
||
| This module provides: | ||
| - torch_cached_mla_with_cache: cached backend op | ||
| - TorchBackendMLAAttention: attention descriptor | ||
|
|
||
| FlashInfer MLA Cache Layout: | ||
| mla_cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] | ||
| - No num_heads dimension (MLA-specific optimization) | ||
| - compressed_kv_cached = mla_cache[:, :, :kv_lora_rank] (zero-copy slice) | ||
| - kpe_cached = mla_cache[:, :, kv_lora_rank:] (zero-copy slice) | ||
|
|
||
| The implementation uses: | ||
| - Prefill: Expand compressed_kv -> full K, V, compute normal attention | ||
| - Generate: Weight absorption for efficiency (Q @ W^T instead of expanding cached KV) | ||
|
|
||
| Reference: https://docs.flashinfer.ai/tutorials/kv_layout.html#mla-page-layout | ||
| """ | ||
|
|
||
| import math | ||
| from typing import List, Optional | ||
|
|
||
| import torch |
There was a problem hiding this comment.
Missing NVIDIA copyright header.
As per coding guidelines, all .py source files should contain an NVIDIA copyright header.
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py` around
lines 1 - 23, Add the required NVIDIA copyright header at the top of this module
file (before the module docstring) per project guidelines; update the file
containing torch_cached_mla_with_cache and TorchBackendMLAAttention to prepend
the standard NVIDIA comment block (including copyright year(s) and "NVIDIA
CORPORATION" wording, license reference if required by repo policy) so every .py
source file includes the header.
| # Store Q output (CONTIGUOUS layout) | ||
| # Output layout: [B, S, H_Q, D] with first half = y1, second half = y2 | ||
| # Compute contiguous strides: stride_b=S*H_Q*D, stride_s=H_Q*D, stride_h=D, stride_d=1 | ||
| q_out_stride_b = S * H_Q * D | ||
| q_out_stride_s = H_Q * D | ||
| q_out_stride_h = D | ||
| q_out_offsets_first = ( | ||
| batch_id * q_out_stride_b | ||
| + seq_offsets[:, None, None] * q_out_stride_s | ||
| + head_offsets[None, :, None] * q_out_stride_h | ||
| + dim_offsets[None, None, :] # stride_d = 1 | ||
| ) | ||
| q_out_offsets_second = q_out_offsets_first + D2 # D2 * 1 | ||
|
|
||
| tl.store(q_out_ptr + q_out_offsets_first, q_y1, mask=load_mask) | ||
| tl.store(q_out_ptr + q_out_offsets_second, q_y2, mask=load_mask) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Step 1: Find the caller function and examine it
echo "=== Looking for triton_rope_on_interleaved_qk_inputs function and its usage ==="
fd "triton_rope" --type=py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 329
🏁 Script executed:
#!/bin/bash
# Step 2: Find all usages of the function
rg "triton_rope_on_interleaved_qk_inputs" --type=py -B 3 -A 3Repository: NVIDIA/TensorRT-LLM
Length of output: 7891
🏁 Script executed:
#!/bin/bash
# Step 3: Examine the triton_rope.py file around the allocation
find . -name "triton_rope.py" -type f | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 127
🏁 Script executed:
#!/bin/bash
# Step 4: Check for torch.empty_like usage and any contiguity handling
rg "empty_like" --type=py -B 2 -A 2 | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 3488
🏁 Script executed:
#!/bin/bash
# Step 5: Examine the triton_rope.py file to see the actual implementation
cat tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py | head -250Repository: NVIDIA/TensorRT-LLM
Length of output: 7730
🏁 Script executed:
#!/bin/bash
# Step 6: Check the validation function in rope.py to see if contiguity is validated
rg "_validate_interleaved_rope_inputs|is_contiguous" . -A 5 -B 2Repository: NVIDIA/TensorRT-LLM
Length of output: 47644
🏁 Script executed:
#!/bin/bash
# Step 7: Look at the test to see if inputs are contiguous
cat tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py | head -150Repository: NVIDIA/TensorRT-LLM
Length of output: 6046
🏁 Script executed:
#!/bin/bash
# Step 8: Check for any contiguity validation/enforcement in the caller
cat tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py | sed -n '141,220p'Repository: NVIDIA/TensorRT-LLM
Length of output: 2780
🏁 Script executed:
#!/bin/bash
# Step 9: Look at the kernel implementation to confirm hardcoded stride assumptions
cat tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_kernels/rope.py | sed -n '146,165p'Repository: NVIDIA/TensorRT-LLM
Length of output: 165
🏁 Script executed:
#!/bin/bash
# Step 10: Check if there are any other tests that create non-contiguous tensors
rg "transpose|permute|view.*-1" tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 9356
🏁 Script executed:
#!/bin/bash
# Step 11: Locate the kernel file and examine it
find . -name "triton_rope_kernel.py" -type fRepository: NVIDIA/TensorRT-LLM
Length of output: 134
🏁 Script executed:
#!/bin/bash
# Step 12: Check the kernel implementation around line 146-161
fd "triton_rope_kernel" -x cat {} \; | sed -n '140,170p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1234
🏁 Script executed:
#!/bin/bash
# Step 13: Verify torch.empty_like behavior with memory format
python3 << 'EOF'
import torch
# Test 1: empty_like with contiguous tensor
x_contig = torch.randn(2, 3, 4, 5)
y_contig = torch.empty_like(x_contig)
print(f"Contiguous input: is_contiguous={x_contig.is_contiguous()}, strides={x_contig.stride()}")
print(f"empty_like output: is_contiguous={y_contig.is_contiguous()}, strides={y_contig.stride()}")
print()
# Test 2: empty_like with non-contiguous tensor
x_noncontig = x_contig.transpose(0, 1)
y_noncontig = torch.empty_like(x_noncontig)
print(f"Non-contiguous input: is_contiguous={x_noncontig.is_contiguous()}, strides={x_noncontig.stride()}")
print(f"empty_like output: is_contiguous={y_noncontig.is_contiguous()}, strides={y_noncontig.stride()}")
EOFRepository: NVIDIA/TensorRT-LLM
Length of output: 180
🏁 Script executed:
#!/bin/bash
# Step 14: Check whether inputs to triton_rope_on_interleaved_qk_inputs are validated for contiguity
rg "contiguous|is_contiguous" tensorrt_llm/_torch/auto_deploy/custom_ops/rope/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 2658
🏁 Script executed:
#!/bin/bash
# Step 15: Check the complete signature and context for output allocation
cat tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py | sed -n '155,210p'Repository: NVIDIA/TensorRT-LLM
Length of output: 2251
🏁 Script executed:
#!/bin/bash
# Step 16: Look for any contiguity enforcement in upstream callers
rg "triton_rope_on_interleaved_qk_inputs|apply_rope_on_interleaved_qk_inputs" tensorrt_llm/_torch/auto_deploy/transform/library/ -B 5 -A 5Repository: NVIDIA/TensorRT-LLM
Length of output: 3365
Kernel output strides assume contiguous layout, but torch.empty_like() preserves input memory format.
The kernel computes output strides assuming perfectly contiguous [B, S, H_Q, D] and [B, S, H_K, D] layouts (e.g., q_out_stride_b = S * H_Q * D). However, the caller allocates outputs with torch.empty_like(q), which preserves the memory layout of the input tensor via torch.preserve_format. If q or k are non-contiguous (e.g., after a transpose or slice), the output tensors inherit that non-contiguity, causing the kernel to compute offsets that write to incorrect memory locations.
No validation enforces that inputs are contiguous, making this a silent correctness bug when non-contiguous tensors are passed.
Fix: Use explicit contiguous allocation
- q_out = torch.empty_like(q)
- k_out = torch.empty_like(k)
+ q_out = torch.empty(q.shape, dtype=q.dtype, device=q.device)
+ k_out = torch.empty(k.shape, dtype=k.dtype, device=k.device)🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope_kernel.py` around
lines 146 - 161, The kernel computes q/k output offsets assuming contiguous
[B,S,H,D] layout (see q_out_stride_b, q_out_stride_s, q_out_stride_h,
q_out_offsets_first and the tl.store calls to q_out_ptr), but outputs are
allocated with torch.empty_like which may preserve a non-contiguous memory
format; fix by ensuring inputs and outputs are contiguous before launching the
Triton kernel: make q and k contiguous (or allocate outputs explicitly with
memory_format=torch.contiguous_format) in the wrapper that prepares
q_out_ptr/k_out_ptr so the calculated strides/offsets match the actual memory
layout used by tl.store.
| @@ -0,0 +1,655 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
Copyright year should be 2026.
Same as the GLM4 file — newly created in February 2026 but header says 2025. As per coding guidelines, source files should contain the year of latest meaningful modification.
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py` at line
1, Update the copyright header in the modeling_deepseek.py source file from 2025
to 2026; locate the top-of-file copyright comment in
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py and change
the year to 2026 so the header reflects the latest meaningful modification year.
| @@ -0,0 +1,830 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
Copyright year should be 2026.
The file is newly created in February 2026, but the copyright header says 2025. As per coding guidelines, source files should contain an NVIDIA copyright header with the year of latest meaningful modification.
Proposed fix
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.📝 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.
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. |
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py` at
line 1, Update the file header comment that currently reads "Copyright (c) 2025,
NVIDIA CORPORATION. All rights reserved." to reflect the latest modification
year 2026; replace "2025" with "2026" in the top-of-file copyright string in
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py so the
header becomes "Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved."
| @pytest.mark.parametrize("batch_size", [4]) | ||
| @pytest.mark.parametrize("dtype", [torch.bfloat16]) | ||
| @pytest.mark.parametrize("device", ["cuda"]) | ||
| def test_flashinfer_mla_init_decode_wrapper_with_buffers(batch_size, dtype, device): | ||
| """Test that _init_decode_wrapper correctly passes buffer tensors with use_cuda_graph=True. | ||
|
|
||
| This test directly tests the _init_decode_wrapper method to verify buffer handling. | ||
| """ | ||
| # Reset planner | ||
| _GlobalFlashInferMLAPlanner.workspace_buffer = None | ||
| _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} | ||
| _GlobalFlashInferMLAPlanner.reset(torch.device(device)) |
There was a problem hiding this comment.
Unused dtype parameter in test.
The dtype parameter is parameterized (Line 2135) but never used in the test body — all buffer tensors are created with torch.int32. Either remove the dtype parametrization or use it when creating relevant tensors.
Proposed fix (remove unused parametrize)
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("device", ["cuda"])
-def test_flashinfer_mla_init_decode_wrapper_with_buffers(batch_size, dtype, device):
+@pytest.mark.parametrize("device", ["cuda"])
+def test_flashinfer_mla_init_decode_wrapper_with_buffers(batch_size, device):📝 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.
| @pytest.mark.parametrize("batch_size", [4]) | |
| @pytest.mark.parametrize("dtype", [torch.bfloat16]) | |
| @pytest.mark.parametrize("device", ["cuda"]) | |
| def test_flashinfer_mla_init_decode_wrapper_with_buffers(batch_size, dtype, device): | |
| """Test that _init_decode_wrapper correctly passes buffer tensors with use_cuda_graph=True. | |
| This test directly tests the _init_decode_wrapper method to verify buffer handling. | |
| """ | |
| # Reset planner | |
| _GlobalFlashInferMLAPlanner.workspace_buffer = None | |
| _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} | |
| _GlobalFlashInferMLAPlanner.reset(torch.device(device)) | |
| `@pytest.mark.parametrize`("batch_size", [4]) | |
| `@pytest.mark.parametrize`("device", ["cuda"]) | |
| def test_flashinfer_mla_init_decode_wrapper_with_buffers(batch_size, device): | |
| """Test that _init_decode_wrapper correctly passes buffer tensors with use_cuda_graph=True. | |
| This test directly tests the _init_decode_wrapper method to verify buffer handling. | |
| """ | |
| # Reset planner | |
| _GlobalFlashInferMLAPlanner.workspace_buffer = None | |
| _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} | |
| _GlobalFlashInferMLAPlanner.reset(torch.device(device)) |
🧰 Tools
🪛 Ruff (0.14.14)
[warning] 2137-2137: Unused function argument: dtype
(ARG001)
🤖 Prompt for AI Agents
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py`
around lines 2134 - 2145, The test
test_flashinfer_mla_init_decode_wrapper_with_buffers declares a parameterized
dtype but never uses it; either remove the `@pytest.mark.parametrize`("dtype",
[torch.bfloat16]) line or apply the dtype when constructing buffer tensors in
the test so the parameter is exercised (inspect the tensor creation sites in
test_flashinfer_mla_init_decode_wrapper_with_buffers and update them to use the
parameter or delete the unused parametrize), referencing the test function name
and the _GlobalFlashInferMLAPlanner setup to locate the code.
| scale: float = None, | ||
| kv_lora_rank: int = None, |
There was a problem hiding this comment.
Use explicit Optional type hints instead of implicit None defaults.
PEP 484 prohibits implicit Optional. The type hints scale: float = None and kv_lora_rank: int = None should use Optional[float] and Optional[int] respectively. This is also flagged by Ruff (RUF013).
Proposed fix
Add Optional to imports at the top of the file:
from typing import OptionalThen:
- scale: float = None,
- kv_lora_rank: int = None,
+ scale: Optional[float] = None,
+ kv_lora_rank: Optional[int] = None,🧰 Tools
🪛 Ruff (0.14.14)
[warning] 33-33: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
[warning] 34-34: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
🤖 Prompt for AI Agents
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_torch_mla_op.py`
around lines 33 - 34, Add Optional to the typing imports and change the two
parameters that default to None to explicit Optional types: replace the implicit
"scale: float = None" with "scale: Optional[float]" and "kv_lora_rank: int =
None" with "kv_lora_rank: Optional[int]" in the function or signature where
these parameters are declared (look for the parameters named scale and
kv_lora_rank in the test_torch_mla_op.py signature) so the type hints comply
with PEP 484 and Ruff RUF013.
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
This reverts commit 5ead4b4.
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
models are uploaded to llm-models, tests can be unwaived Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
|
all prs already merged to main |
4 participants
Summary by CodeRabbit
Description
This is a bundle of the three PRs (running CI and getting merged). The bundle can be used to try out the model with latest perf enhancements in the meanwhile.
#11324
#11351
#11327
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.