fix(collector): align model dtype for vLLM FlashInfer fp8 KV attention

[kaim-eng]

Summary

• collector/vllm/collect_attn.py: When fp8 KV cache is enabled and vLLM routes attention to FLASHINFER, align vllm_config.model_config.dtype with kv_cache_spec.dtype before FlashInferMetadataBuilder init. vLLM 0.22 asserts these dtypes match; the collector's KV spec uses fp8 while model_config stayed bfloat16, causing AssertionError on every fp8+head_dim=128 case on SM100.

Context

Discovered during B200 vLLM 0.22.0 power re-collection (R2). fp8+head_dim=192 used a different backend and did not hit this path. Targeted re-collect with this WAR produced 11,288 context + 11,333 generation rows (job 3137493, umb-b200-236).

Independent of #1358 (helper + TRT-LLM MLA API) and #1359 (TRT-LLM kernel-limit FIXME comments).

Test plan

• pytest tests/unit/collector/test_vllm_collect_attn.py -m unit
• Smoke: collector/collect.py --backend vllm --ops attention_generation --case-filter "128, True" --smoke on vLLM 0.22.x + SM100 — confirm fp8+head_dim=128 cases run without FlashInfer dtype assert

Summary by CodeRabbit

• Bug Fixes
  • Fixed FlashInfer attention processing with FP8 KV caches by ensuring consistent data types during metadata setup.
  • Prevented dtype mismatch errors when the model configuration uses a different default type.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b58512a4-1e10-49f9-b599-338f6317e676

📥 Commits

Reviewing files that changed from the base of the PR and between 92b4ce8 and b7dd443.

📒 Files selected for processing (1)

• collector/vllm/collect_attn.py

📜 Recent 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 (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/vllm/collect_attn.py

collector/**

📄 CodeRabbit inference engine (.claude/rules/generator/cross_module_impact.md)

Review the collector when generated configuration formats change or when generator parameter names used in benchmark configurations change.

collector/**: Keep collector-layer rules in their permitted locations: base operation YAML for sweep parameters and cited quant-mode SM limits; model case YAML for structural shapes and policies; capabilities YAML for positive hardware dtype/op minimum-SM floors; backend registries for version routes and unverified markers; collectors for dispatch, classified runtime errors, and memory-feasibility generation filters; and the denylist only for hangs or node-killing cases.
Do not silently skip queued cases in collector code. A queued case must be executed or produce a classified exception; only generation-time capability floors, registry unverified markers, and the hang denylist may prevent execution.
Use framework-selected kernel/backend dispatch for measurements. Prefer the framework builder or selector; manually pin a backend only with a pinned-version file-and-line source citation; record the actually invoked kernel in kernel_source; and re-verify pins and FIXME(kernel-limit) notes on framework version bumps.
Do not invent backend fallbacks or swap backends in exception handlers. If the framework-selected path cannot be constructed or invoked, raise a classified error. Replicate fallback chains only when the framework itself uses them, and record the actual path.
The only sanctioned in-collector filter is memory feasibility: perform it only during generation inside get_*_test_cases(), use size-versus-capacity arithmetic based on live device memory or an approved capacity fallback, never use framework-version or model-name predicates, and log dropped counts in the required format.
Keep framework kernel limits as FIXME(kernel-limit) comments at the owning collector invocation site, including the claimed limit, origin, and unverified status. Do not encode unverified limits in YAML; affected cases should fail at runtime...

Files:

• collector/vllm/collect_attn.py

---

⚙️ CodeRabbit configuration file

collector/**: - Enforce the collector rules from .claude/rules/collector/layer_permissions.md, failure_handling.md, and case_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.yaml is for hang/node-killers only, dated.
• Collector changes must stay within collector/ and tests/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/vllm/collect_attn.py

collector/**/*.py

📄 CodeRabbit inference engine (.claude/rules/collector/case_authoring.md)

collector/**/*.py: If a generator lacks support for a required model-correlated field, extend the generator mechanism and obtain human review; do not approximate the behavior with post-generation filtering.
Unresolvable declared inputs must raise; generators must not substitute defaults, another model's geometry, or a close-enough quantization mode. Planned operations that expand to zero cases must have logged drop reasons, otherwise the population bug must be fixed.

Files:

• collector/vllm/collect_attn.py

**/*

📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)

**/*: Treat this file as the only always-injected rule file; do not add new always-on rules here without human approval. New rule files must include paths: frontmatter.
When reviewing changes in a governed area, read that area's rule files first, even if path-based auto-loading does not occur during a read-only review. Rule violations are review findings even when the code works.
A task must remain within its module: collector tasks may change only collector/ and its tests, generator tasks only src/aiconfigurator/generator/ and its tests, and SDK tasks must not modify either. Cross-module contract changes require explicit human approval.

Files:

• collector/vllm/collect_attn.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/vllm/collect_attn.py

collector/**/*

📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)

Changes under collector/ must follow .claude/rules/collector/layer_permissions.md and failure_handling.md; case YAML work must also follow case_authoring.md. Do not apply generator rules to collector/case_generator.py.

Files:

• collector/vllm/collect_attn.py

**/*.{py,pyi}

📄 CodeRabbit inference engine (AGENTS.md)

Run ruff check . and ruff format --check . to validate Python linting and formatting.

Files:

• collector/vllm/collect_attn.py

🧠 Learnings (3)

📚 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/vllm/collect_attn.py

📚 Learning: 2026-07-05T07:16:31.423Z

Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 1264
File: collector/vllm/utils.py:29-34
Timestamp: 2026-07-05T07:16:31.423Z
Learning: In ai-dynamo/aiconfigurator’s vLLM collector modules, keep vLLM imports at top level (not lazily deferred) so a mismatched vLLM version fails during the per-op import and is recorded/classified as a per-op failure (via the existing per-op try/except and __compat gating), rather than altering failure order/granularity. Do not rely on moving vLLM imports into shared helpers like collector/vllm/utils.py to change this behavior. Exception: collector/vllm/collect_mhc_module.py is intentionally “import-vllm-free” (vLLM dependencies were localized) and should remain so to preserve the intended behavior there.

Applied to files:

• collector/vllm/collect_attn.py

📚 Learning: 2026-07-11T03:07:50.673Z

Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 1344
File: tests/unit/collector/test_framework_manifest.py:42-49
Timestamp: 2026-07-11T03:07:50.673Z
Learning: For the vLLM “collector v2” registries (e.g., collector/vllm/registry.py and any other vLLM collector v2 registry modules), treat the registry entry `versions` field as a retired per-entry versioning mechanism. Code and tests should expect it to be empty/falsy for every registry entry (i.e., review changes to ensure `entry.versions` is not populated). Version pinning for the vLLM backend must be centralized in `collector/framework_manifest.yaml` via `get_collector_runtime("vllm")`, and each vLLM collector module must declare exactly one `__compat__ = "vllm==<manifest_version>"` line matching the pinned manifest version. If relevant, do not misread the framework-manifest test: `test_active_cuda_vllm_collectors_are_exactly_pinned_to_manifest_version` asserts `all(not entry.versions for entry in VLLM_REGISTRY)` and that `__compat__` matches the manifest version (it is not an inverted assertion).

Applied to files:

• collector/vllm/collect_attn.py

🔇 Additional comments (1)

collector/vllm/collect_attn.py (1)

210-213: LGTM!

---

Walkthrough

Updates the FlashInfer attention path so FP8 KV-cache execution synchronizes model_config.dtype with kv_cache_spec.dtype before metadata construction.

Changes

FlashInfer dtype alignment

Layer / File(s)	Summary
Align model and KV-cache dtypes collector/vllm/collect_attn.py	When FP8 KV caching is enabled, run_attention_torch assigns kv_cache_spec.dtype to vllm_config.model_config.dtype in the FlashInfer path.

Estimated code review effort: 1 (Trivial) | ~3 minutes

Poem

FlashInfer checks the dtype in flight,
FP8 caches now match just right.
Model and KV walk side by side,
Metadata flows with nothing to hide.

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the vLLM FlashInfer fp8 KV dtype alignment fix.
Description check	✅ Passed	The description covers the change, context, and test plan, but it omits several template sections like reviewer start and related issues.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}