fix(collector): CPU tensor perf regression and TRT-LLM KV-cache API migration#1358
fix(collector): CPU tensor perf regression and TRT-LLM KV-cache API migration#1358kaim-eng wants to merge 5 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe changes move MoE redistribution calculations to CPU-backed integer tensors, simplify rank-load selection, add focused helper tests, and support batched KV-cache sequence insertion with fallback compatibility. ChangesMoE Distribution Updates
Batched KV-Cache Setup
Estimated code review effort: 3 (Moderate) | ~20 minutes Poem
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
collector/helper.py (1)
1073-1095: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd regression coverage for device preservation and exact totals.
Cover CPU input, CUDA/default-device input when available, overflow redistribution, positive and negative deltas, returned-device preservation, and exact target sums. A one-click suggestion is not safe because the relevant test fixtures are outside the supplied diff.
Also applies to: 1126-1131, 1178-1179, 1239-1250, 1259-1266
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@collector/helper.py` around lines 1073 - 1095, Add regression tests for the redistribution logic covering CPU input and CUDA/default-device input when available, including positive and negative overflow deltas. Assert the returned tensor preserves the original device and that each result matches the exact target totals, including large overflow cases; cover the related redistribution paths around the visible counts_2d logic and associated call sites.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@collector/helper.py`:
- Around line 1073-1095: Add regression tests for the redistribution logic
covering CPU input and CUDA/default-device input when available, including
positive and negative overflow deltas. Assert the returned tensor preserves the
original device and that each result matches the exact target totals, including
large overflow cases; cover the related redistribution paths around the visible
counts_2d logic and associated call sites.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1b9e5a66-f990-443c-bc62-16abb7d8a1c5
📒 Files selected for processing (2)
collector/helper.pycollector/trtllm/collect_mla.py
📜 Review details
⏰ Context from checks skipped due to timeout. (6)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Build and Test (unit)
- GitHub Check: Cargo Deny
- GitHub Check: Rust/Python engine-step parity
- GitHub Check: Build and Test (e2e)
🧰 Additional context used
📓 Path-based instructions (6)
{collector/**,tests/unit/collector/**}
📄 CodeRabbit inference engine (.claude/rules/collector/failure_handling.md)
{collector/**,tests/unit/collector/**}: In collector failure handling, observe failures rather than predict them: do not use a declarative expected-failure layer or automatic skips; failing groups must be fixed.
Record every worker failure automatically in the relevant error and collection-summary JSON files with classification, case parameters, exception details, and model/dtype grouping.
Reset a worker after CUDA-fatal errors only after recording the failed task.
Treat missing data points as tolerable downstream because the SDK can interpolate, extrapolate, reuse sibling-version rows, or use HYBRID empirical estimates.
For a hanging case or one that kills the node, add a dated denylist.yaml entry with a reason.
For an entirely unvalidated operation/backend pair, mark the registry OpEntry as unverified=true.
For cases not validated on a specific SM, mark the registry OpEntry with unverified_sms=(sm,).
Represent physically impossible hardware capabilities with a positive floor in capabilities.yaml; do not use it for framework-version kernel gaps.
Treat out-of-memory failures as unclassified until reproduced on a clean GPU; only a genuine OOM may be excluded by the sanctioned generation-time memory filter.
Fix proven collector-code bugs in code rather than using skips, and re-check dispatch and skip rules.
Do not encode framework-version gaps in YAML; allow them to fail and be re-tested after version changes.
Treat isolated, explained, unclustered failures as acceptable; investigate around 10% unexpected failures or sooner when failures cluster.
Treat roughly one-third failing cases or an entire failing family as a systemic collector problem; stop collection and fix the collector.
Never hide failures with broad skips, retries, generic OOM labels, reduced coverage, synthetic rows, or weakened benchmarks.
Compare failure records with the previous run for the same backend and version; prioritize only new failure groups.
Do not mechanically translate failure records in...
Files:
collector/trtllm/collect_mla.pycollector/helper.py
collector/**
📄 CodeRabbit inference engine (.claude/rules/collector/layer_permissions.md)
collector/**: Keep collector-layer rules in their designated locations: base-op YAML for sweeps and axis-levelmin_sm; model case YAML for structural shapes and activation; capabilities YAML for positive hardware floors; registries for version routing and maturity markers; collectors for dispatch, classified errors, and memory-feasibility filtering; and the denylist only for hangs or node-killing cases.
Collector tasks may modify onlycollector/andtests/unit/collector/; they must not modify SDK, Rust, tools, systems data, or generator code.
Changing the collector data contract requires explicit human approval; do not independently add or rename columns, perf files, or SDK-parsed key dimensions.
Do not edit human-owned collector rule files under.claude/rules/collector/as a side effect of a fix task; propose policy changes instead.
For a failing case, default to no code change: verify that it is recorded and classified in the failure log, then stop unless the failure-handling decision tree requires escalation.
Mechanism changes to the executor, case generator, failure classification, or case-ID format require explicit human approval and must be proposed rather than included in a case fix.Review the collector when generated configuration formats change or when generator parameter names used in benchmark configurations change.
Files:
collector/trtllm/collect_mla.pycollector/helper.py
⚙️ CodeRabbit configuration file
collector/**: - Check collector changes for backend/runtime version accuracy, GPU resource assumptions, reproducible command construction, and clear failure evidence.
- Flag changes that make support-matrix or perf-data results harder to trace back to the command, model, system, quantization, or runtime version that produced them.
Files:
collector/trtllm/collect_mla.pycollector/helper.py
collector/**/collect_*.py
📄 CodeRabbit inference engine (.claude/rules/collector/layer_permissions.md)
collector/**/collect_*.py: Collector code may dispatch between kernel paths by SM or version, record the actualkernel_source, raise classified exceptions for runtime probes, and perform memory-feasibility filtering; it must not silently skip queued cases, perform other case filtering, or patch code for one shape.
A queued case must be executed or raise a classified error; branches may change how a case runs, but must not change whether it runs.
Use the framework's own dispatch path to select kernels. Manual backend pinning requires a source citation, andkernel_sourcemust record the actually invoked kernel.
Do not invent backend fallbacks or swap backends in exception handlers. Raise a classified error if the framework-selected path cannot run; only replicate fallback behavior performed by the framework itself.
Measurement-method degradation, such as CUDA-graph capture falling back to eager execution, is allowed only when recorded in the output row; API-compatibility shims may alter construction but never kernel selection.
Memory-feasibility filtering is permitted only during generation insideget_*_test_cases(), using footprint-versus-capacity arithmetic and live device memory where possible; drops must be counted and logged, never silently skipped.
Framework kernel limits must remain asFIXME(kernel-limit)comments at the owning collector invocation site, stating the limit, origin, and unverified status; do not encode them in YAML or implement guards from unverified claims.
On framework version bumps, re-verify every pinned backend and everyFIXME(kernel-limit)against framework source before trusting collected data.
Files:
collector/trtllm/collect_mla.py
collector/**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
collector/**/*: When editing or reviewingcollector/**, load and applylayer_permissions.mdandfailure_handling.md; also loadcase_authoring.mdfor case YAML work.
Do not apply generator-module rules tocollector/case_generator.py; it expands collection test cases and is unrelated to deployment-config generation.
Collector work is standalone GPU performance data collection and is not part of the wheel runtime.
Files:
collector/trtllm/collect_mla.pycollector/helper.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
collector/trtllm/collect_mla.pycollector/helper.py
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
collector/trtllm/collect_mla.pycollector/helper.py
🧠 Learnings (1)
📚 Learning: 2026-02-28T11:44:28.109Z
Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 466
File: collector/trtllm/collect_moe_v3.py:4-4
Timestamp: 2026-02-28T11:44:28.109Z
Learning: In collector-related Python files (e.g., collector/trtllm/collect_moe_v3.py), document and enforce that version incompatibilities are surfaced as non-fatal runtime errors during collection. Do not add aggressive preventive version-gating; instead, allow generating test cases that may not be supported across all versions within the __compat__ range and rely on runtime error handling to skip or flag unsupported cases. This should be verifiable by ensuring collection proceeds, errors are reported, and unsupported cases do not halt the overall process.
Applied to files:
collector/trtllm/collect_mla.pycollector/helper.py
🔇 Additional comments (3)
collector/helper.py (2)
1073-1095: LGTM!Also applies to: 1126-1131, 1178-1179, 1239-1250, 1259-1266
1073-1095: 📐 Maintainability & Code QualityRun the Ruff checks in the uv environment before merge
uv run ruff check .anduv run ruff format --check .collector/trtllm/collect_mla.py (1)
358-360: LGTM!Also applies to: 370-379
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/collector/test_helper_moe_distribution.py`:
- Around line 89-91: Extend test_power_law_distribution_exact_sum to assert that
EP rank 0 receives the maximum load from _generate_power_law_distribution.
Preserve the existing global total assertion and compare rank 0’s count against
the maximum count across EP ranks.
- Around line 16-19: Add the project’s unit-test marker to the test module in
tests/unit/collector/test_helper_moe_distribution.py, applying it at module
scope so all tests in the file are selected by pytest -m unit. Reuse the
existing pytest marker convention used by nearby unit-test modules.
- Around line 16-19: Fix the Ruff I001 import-order violation in
test_helper_moe_distribution.py by applying the repository’s expected import
grouping and sorting for pytest, torch, and the collector.helper imports. Keep
the imported symbols unchanged and limit the change to the import block.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c47a0b89-de14-4c26-81d7-b72e06aadd2d
📒 Files selected for processing (1)
tests/unit/collector/test_helper_moe_distribution.py
📜 Review details
⏰ Context from checks skipped due to timeout. (6)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Cargo Deny
- GitHub Check: Rust/Python engine-step parity
- GitHub Check: Build and Test (unit)
- GitHub Check: Build and Test (e2e)
⚠️ CI failures not shown inline (1)
GitHub Actions: Lint and Format / 1_Lint and Format (Ruff).txt: fix(collector): CPU tensor perf regression and TRT-LLM KV-cache API migration
Conclusion: failure
##[group]Run ruff check .
�[36;1mruff check .�[0m
shell: /usr/bin/bash -e {0}
env:
pythonLocation: /opt/hostedtoolcache/Python/3.12.13/x64
PKG_CONFIG_PATH: /opt/hostedtoolcache/Python/3.12.13/x64/lib/pkgconfig
Python_ROOT_DIR: /opt/hostedtoolcache/Python/3.12.13/x64
Python2_ROOT_DIR: /opt/hostedtoolcache/Python/3.12.13/x64
Python3_ROOT_DIR: /opt/hostedtoolcache/Python/3.12.13/x64
LD_LIBRARY_PATH: /opt/hostedtoolcache/Python/3.12.13/x64/lib
##[endgroup]
I001 [*] Import block is un-sorted or un-formatted
--> tests/unit/collector/test_helper_moe_distribution.py:16:1
|
14 | """
15 |
16 | / import pytest
17 | | import torch
18 | |
19 | | from collector.helper import _generate_power_law_distribution, _round_robin_adjust_per_rank
| |___________________________________________________________________________________________^
|
help: Organize imports
Found 1 error.
[*] 1 fixable with the `--fix` option.
##[error]Process completed with exit code 1.
🧰 Additional context used
📓 Path-based instructions (7)
{collector/**,tests/unit/collector/**}
📄 CodeRabbit inference engine (.claude/rules/collector/failure_handling.md)
{collector/**,tests/unit/collector/**}: In collector failure handling, observe failures rather than predict them: do not use a declarative expected-failure layer or automatic skips; failing groups must be fixed.
Record every worker failure automatically in the relevant error and collection-summary JSON files with classification, case parameters, exception details, and model/dtype grouping.
Reset a worker after CUDA-fatal errors only after recording the failed task.
Treat missing data points as tolerable downstream because the SDK can interpolate, extrapolate, reuse sibling-version rows, or use HYBRID empirical estimates.
For a hanging case or one that kills the node, add a dated denylist.yaml entry with a reason.
For an entirely unvalidated operation/backend pair, mark the registry OpEntry as unverified=true.
For cases not validated on a specific SM, mark the registry OpEntry with unverified_sms=(sm,).
Represent physically impossible hardware capabilities with a positive floor in capabilities.yaml; do not use it for framework-version kernel gaps.
Treat out-of-memory failures as unclassified until reproduced on a clean GPU; only a genuine OOM may be excluded by the sanctioned generation-time memory filter.
Fix proven collector-code bugs in code rather than using skips, and re-check dispatch and skip rules.
Do not encode framework-version gaps in YAML; allow them to fail and be re-tested after version changes.
Treat isolated, explained, unclustered failures as acceptable; investigate around 10% unexpected failures or sooner when failures cluster.
Treat roughly one-third failing cases or an entire failing family as a systemic collector problem; stop collection and fix the collector.
Never hide failures with broad skips, retries, generic OOM labels, reduced coverage, synthetic rows, or weakened benchmarks.
Compare failure records with the previous run for the same backend and version; prioritize only new failure groups.
Do not mechanically translate failure records in...
Files:
tests/unit/collector/test_helper_moe_distribution.py
tests/unit/collector/**/*.py
📄 CodeRabbit inference engine (.claude/rules/collector/layer_permissions.md)
Mark every new collector unit test with
pytest.mark.unit; otherwise tests that work only with-m unitare invisible to CI.
Files:
tests/unit/collector/test_helper_moe_distribution.py
tests/**/*.{py,yaml,txt,sh}
📄 CodeRabbit inference engine (.claude/rules/generator/testing.md)
Use integration tests for the full input-to-artifacts pipeline, comparing output against golden snapshots without external dependencies.
Files:
tests/unit/collector/test_helper_moe_distribution.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
tests/unit/collector/test_helper_moe_distribution.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Run unit tests with
pytest -m unit; run the PR build subset withpytest -m "unit or build"when applicable.
Files:
tests/unit/collector/test_helper_moe_distribution.py
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
tests/unit/collector/test_helper_moe_distribution.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/collector/test_helper_moe_distribution.py
🧠 Learnings (1)
📚 Learning: 2026-02-28T11:44:28.109Z
Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 466
File: collector/trtllm/collect_moe_v3.py:4-4
Timestamp: 2026-02-28T11:44:28.109Z
Learning: In collector-related Python files (e.g., collector/trtllm/collect_moe_v3.py), document and enforce that version incompatibilities are surfaced as non-fatal runtime errors during collection. Do not add aggressive preventive version-gating; instead, allow generating test cases that may not be supported across all versions within the __compat__ range and rely on runtime error handling to skip or flag unsupported cases. This should be verifiable by ensuring collection proceeds, errors are reported, and unsupported cases do not halt the overall process.
Applied to files:
tests/unit/collector/test_helper_moe_distribution.py
🪛 GitHub Actions: Lint and Format / 1_Lint and Format (Ruff).txt
tests/unit/collector/test_helper_moe_distribution.py
[error] 16-19: ruff check failed (I001): Import block is un-sorted or un-formatted. Organize imports (ruff can auto-fix with --fix).
🪛 GitHub Actions: Lint and Format / Lint and Format (Ruff)
tests/unit/collector/test_helper_moe_distribution.py
[error] 16-19: Ruff check failed: import block is un-sorted or un-formatted. Rule: I001. Help: Organize imports (1 fixable with --fix). Command: ruff check .
…igration helper.py: move count tensor to CPU during round-robin redistribution to avoid per-iteration GPU syncs when torch.set_default_device(cuda) is active in collector workers. Replace Conv1d rank-sum with view().sum() for the same reason. Add int() casts where .item() can return float to prevent silent type errors in range() and arithmetic. collect_mla.py: TRT-LLM >=1.3.0rc era removed per-sequence KVCacheManager.impl.add_sequence() in favour of a batched add_sequence_batch(request_infos, requests). Buffer requests across the loop and dispatch via the batched API, falling back to the legacy per-sequence call for older runtimes still covered by __compat__. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Kai Ma <kaim@nvidia.com>
Covers _round_robin_adjust_per_rank (device preservation, exact-total invariant for add/subtract, no-op at zero remaining, early-exit when no valid slot) and _generate_power_law_distribution (sum == num_tokens*topk, per-expert upper bound, output length, assignment shape). The add_sequence_batch shim in collect_mla._run_attn_for_backend is noted in the test file as untested here — it is embedded in a function that takes live TRT-LLM KV-cache objects and would need the dispatch extracted to be unit-testable without mocking the full stack. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Kai Ma <kaim@nvidia.com>
- Add pytestmark = pytest.mark.unit so tests are visible to pytest -m unit - Fix Ruff I001: split multi-symbol import to parenthesised form - Add rank-0 max-load assertion to test_power_law_distribution_exact_sum, covering the view().sum().argmax() path that replaced the Conv1d approach Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Kai Ma <kaim@nvidia.com>
…moe_distribution
Unit CI image does not install torch; bare `import torch` caused
ModuleNotFoundError and blocked collection of the entire test module.
Replace with `torch = pytest.importorskip("torch")` so the module is
skipped cleanly when torch is absent and runs normally when present.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Kai Ma <kaim@nvidia.com>
95d9ce0 to
faa46e3
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
collector/trtllm/collect_mla.py (1)
358-379: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo isolated test coverage for the batched/fallback shim.
The new
test_helper_moe_distribution.pydocstring already flags that this shim isn't unit-tested in isolation because it's embedded in a large function requiring live TRT-LLM objects. Extracting thehasattrdispatch into a small helper (takingkv_cache_manager.impl,batch_request_infos,batch_llm_requests) would let both branches be exercised with a mock object, closing that acknowledged gap.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@collector/trtllm/collect_mla.py` around lines 358 - 379, The KV-cache sequence dispatch is embedded in the larger generation flow and lacks isolated branch coverage. Extract the hasattr-based dispatch around kv_cache_manager.impl into a small helper accepting the implementation, batch_request_infos, and batch_llm_requests; preserve the batched add_sequence_batch path and legacy per-sequence add_sequence fallback, then have the existing flow call this helper.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@collector/trtllm/collect_mla.py`:
- Around line 358-379: The KV-cache sequence dispatch is embedded in the larger
generation flow and lacks isolated branch coverage. Extract the hasattr-based
dispatch around kv_cache_manager.impl into a small helper accepting the
implementation, batch_request_infos, and batch_llm_requests; preserve the
batched add_sequence_batch path and legacy per-sequence add_sequence fallback,
then have the existing flow call this helper.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6641926d-5d3b-4127-a000-9aae890a550a
📒 Files selected for processing (2)
collector/helper.pycollector/trtllm/collect_mla.py
🚧 Files skipped from review as they are similar to previous changes (1)
- collector/helper.py
📜 Review details
⏰ Context from checks skipped due to timeout. (6)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Build and Test (e2e)
- GitHub Check: Build and Test (unit)
- GitHub Check: Cargo Deny
- GitHub Check: Rust/Python engine-step parity
🧰 Additional context used
📓 Path-based instructions (6)
{collector/**,tests/unit/collector/**}
📄 CodeRabbit inference engine (.claude/rules/collector/failure_handling.md)
{collector/**,tests/unit/collector/**}: In collector failure handling, observe failures rather than predict them: do not use a declarative expected-failure layer or automatic skips; failing groups must be fixed.
Record every worker failure automatically in the relevant error and collection-summary JSON files with classification, case parameters, exception details, and model/dtype grouping.
Reset a worker after CUDA-fatal errors only after recording the failed task.
Treat missing data points as tolerable downstream because the SDK can interpolate, extrapolate, reuse sibling-version rows, or use HYBRID empirical estimates.
For a hanging case or one that kills the node, add a dated denylist.yaml entry with a reason.
For an entirely unvalidated operation/backend pair, mark the registry OpEntry as unverified=true.
For cases not validated on a specific SM, mark the registry OpEntry with unverified_sms=(sm,).
Represent physically impossible hardware capabilities with a positive floor in capabilities.yaml; do not use it for framework-version kernel gaps.
Treat out-of-memory failures as unclassified until reproduced on a clean GPU; only a genuine OOM may be excluded by the sanctioned generation-time memory filter.
Fix proven collector-code bugs in code rather than using skips, and re-check dispatch and skip rules.
Do not encode framework-version gaps in YAML; allow them to fail and be re-tested after version changes.
Treat isolated, explained, unclustered failures as acceptable; investigate around 10% unexpected failures or sooner when failures cluster.
Treat roughly one-third failing cases or an entire failing family as a systemic collector problem; stop collection and fix the collector.
Never hide failures with broad skips, retries, generic OOM labels, reduced coverage, synthetic rows, or weakened benchmarks.
Compare failure records with the previous run for the same backend and version; prioritize only new failure groups.
Do not mechanically translate failure records in...
Files:
collector/trtllm/collect_mla.py
collector/**
📄 CodeRabbit inference engine (.claude/rules/collector/layer_permissions.md)
collector/**: Keep collector-layer rules in their designated locations: base-op YAML for sweeps and axis-levelmin_sm; model case YAML for structural shapes and activation; capabilities YAML for positive hardware floors; registries for version routing and maturity markers; collectors for dispatch, classified errors, and memory-feasibility filtering; and the denylist only for hangs or node-killing cases.
Collector tasks may modify onlycollector/andtests/unit/collector/; they must not modify SDK, Rust, tools, systems data, or generator code.
Changing the collector data contract requires explicit human approval; do not independently add or rename columns, perf files, or SDK-parsed key dimensions.
Do not edit human-owned collector rule files under.claude/rules/collector/as a side effect of a fix task; propose policy changes instead.
For a failing case, default to no code change: verify that it is recorded and classified in the failure log, then stop unless the failure-handling decision tree requires escalation.
Mechanism changes to the executor, case generator, failure classification, or case-ID format require explicit human approval and must be proposed rather than included in a case fix.Review the collector when generated configuration formats change or when generator parameter names used in benchmark configurations change.
Files:
collector/trtllm/collect_mla.py
⚙️ CodeRabbit configuration file
collector/**: - Enforce the collector rules from.claude/rules/collector/layer_permissions.md,failure_handling.md, andcase_authoring.md.
- Flag any silent case skip in collector code (a queued case may only execute or raise); the sole sanctioned filter is generation-time memory feasibility with counted drops.
- Flag invented fallbacks on both ends: generation must raise on unresolvable declarations (never substitute defaults or another model's geometry); collectors must never swap in a different backend/kernel than the framework's own dispatch selects — manual pins require framework source citations.
- Flag any reintroduction of selector/exception machinery (case_ids/contains/indices/ranges/limit/rules, sm_exceptions-style shape or version predicates) in YAML or code.
- Capability floors (
cases/capabilities.yaml) may hold hardware facts only: no shapes, no framework versions, no per-backend nesting.cases/denylist.yamlis for hang/node-killers only, dated.- Collector changes must stay within
collector/andtests/unit/collector/; flag producer+consumer contract changes (perf row schema, PerfFile names) unless the PR explicitly declares them.- Check collector changes for backend/runtime version accuracy, GPU resource assumptions, reproducible command construction, and clear failure evidence.
- Flag changes that make support-matrix or perf-data results harder to trace back to the command, model, system, quantization, or runtime version that produced them.
Files:
collector/trtllm/collect_mla.py
collector/**/collect_*.py
📄 CodeRabbit inference engine (.claude/rules/collector/layer_permissions.md)
collector/**/collect_*.py: Collector code may dispatch between kernel paths by SM or version, record the actualkernel_source, raise classified exceptions for runtime probes, and perform memory-feasibility filtering; it must not silently skip queued cases, perform other case filtering, or patch code for one shape.
A queued case must be executed or raise a classified error; branches may change how a case runs, but must not change whether it runs.
Use the framework's own dispatch path to select kernels. Manual backend pinning requires a source citation, andkernel_sourcemust record the actually invoked kernel.
Do not invent backend fallbacks or swap backends in exception handlers. Raise a classified error if the framework-selected path cannot run; only replicate fallback behavior performed by the framework itself.
Measurement-method degradation, such as CUDA-graph capture falling back to eager execution, is allowed only when recorded in the output row; API-compatibility shims may alter construction but never kernel selection.
Memory-feasibility filtering is permitted only during generation insideget_*_test_cases(), using footprint-versus-capacity arithmetic and live device memory where possible; drops must be counted and logged, never silently skipped.
Framework kernel limits must remain asFIXME(kernel-limit)comments at the owning collector invocation site, stating the limit, origin, and unverified status; do not encode them in YAML or implement guards from unverified claims.
On framework version bumps, re-verify every pinned backend and everyFIXME(kernel-limit)against framework source before trusting collected data.
Files:
collector/trtllm/collect_mla.py
collector/**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
collector/**/*: When editing or reviewingcollector/**, load and applylayer_permissions.mdandfailure_handling.md; also loadcase_authoring.mdfor case YAML work.
Do not apply generator-module rules tocollector/case_generator.py; it expands collection test cases and is unrelated to deployment-config generation.
Collector work is standalone GPU performance data collection and is not part of the wheel runtime.
Files:
collector/trtllm/collect_mla.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
collector/trtllm/collect_mla.py
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
collector/trtllm/collect_mla.py
🧠 Learnings (1)
📚 Learning: 2026-02-28T11:44:28.109Z
Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 466
File: collector/trtllm/collect_moe_v3.py:4-4
Timestamp: 2026-02-28T11:44:28.109Z
Learning: In collector-related Python files (e.g., collector/trtllm/collect_moe_v3.py), document and enforce that version incompatibilities are surfaced as non-fatal runtime errors during collection. Do not add aggressive preventive version-gating; instead, allow generating test cases that may not be supported across all versions within the __compat__ range and rely on runtime error handling to skip or flag unsupported cases. This should be verifiable by ensuring collection proceeds, errors are reported, and unsupported cases do not halt the overall process.
Applied to files:
collector/trtllm/collect_mla.py
🔇 Additional comments (1)
collector/trtllm/collect_mla.py (1)
372-379: 🎯 Functional CorrectnessThis batched call matches the TRT-LLM binding signature; keep the fallback for older runtimes.
> Likely an incorrect or invalid review comment.
…im comment Note TRT-LLM version where add_sequence_batch was confirmed present (1.3.0rc15) and the source module to check when retiring the fallback. The hasattr probe remains the live boundary check; the version note tells reviewers when to clean up. Signed-off-by: Kai Ma <kaim@nvidia.com> Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 participant
Summary
collector/helper.py: Move the redistribution count tensor to CPU during the round-robin adjustment loop (_round_robin_adjust_per_rank). Whentorch.set_default_device(cuda)is active in collector workers, each iteration of the tens-of-thousands-step integer loop triggered a GPU sync, making the function prohibitively slow. Also replace theConv1d-based EP-rank sum with a directview().sum()for the same reason. Addint()casts where.item()can return a float to prevent silent type errors inrange()and arithmetic.collector/trtllm/collect_mla.py: TRT-LLM>=1.3.0rcera removed the per-sequenceKVCacheManager.impl.add_sequence()binding in favour of a batchedadd_sequence_batch(request_infos, requests). Buffer all requests across the loop then dispatch via the batched API, with a fallback to the legacy per-sequence call for older runtimes still covered by__compat__.What this is NOT
This PR contains no skip guards and no coverage changes. A second PR will cover the hardware/version-specific skip predicates discovered during the same collection run, with per-guard case-count evidence and runtime repros.
Test plan
torch.set_default_device(cuda)active — confirm no GPU sync hangcollect_mla.pycontext phase against a TRT-LLM>=1.3.0rc15build — confirmadd_sequence_batchpath executes without errorcollect_mla.pycontext phase against a TRT-LLM<1.3.0rcbuild — confirm legacy fallback path executes without error🤖 Generated with Claude Code
Summary by CodeRabbit