[NVBUG-6379624][fix] Enable W4A8 checkpoint loading for Gemma4 K=V layers

[Hudayday]

Description

Gemma4 full-attention layers may set attention_k_eq_v=True. Hugging Face/ModelOpt checkpoints then omit v_proj tensors and reuse k_proj for V.

The Gemma4 HF weight mapper previously synthesized only v_proj.weight. For W4A8_AWQ checkpoints, the missing V quantization tensors left the fused QKV weight_scale width at 18432 while the target expected 20480, causing model
initialization to fail.

This PR:

• Duplicates missing k_proj projection tensors to v_proj for full-attention K=V layers, including W4A8 quantization metadata.
• Preserves explicitly exported V tensors.
• Maps k_scale/k_bias metadata to v_scale/v_bias.
• Supports both text and multimodal checkpoint prefixes without depending on checkpoint dictionary order.
• Adds regression coverage for text/VLM checkpoints, full/sliding attention, quantization metadata, explicit V values, layer-scalar loading, and idempotence.

The change is limited to Gemma4 HF checkpoint preprocessing. It does not change public APIs or inference kernels.

Test Coverage

• Gemma4 weight-mapper regression test:

  pytest -q tests/unittest/_torch/modeling/test_modeling_gemma4.py \
      -k duplicate_full_attention_kv_projection_tensors
  

Passed for both text and multimodal checkpoint layouts.

• Static checks:

  ruff check
  tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
  tests/unittest/_torch/modeling/test_modeling_gemma4.py

  ruff format --check
  tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
  tests/unittest/_torch/modeling/test_modeling_gemma4.py

• B200 end-to-end validation:

  • Exported Gemma4-31B-it using ModelOpt 0.45.0rc1 with W4A8_AWQ and FP8 KV cache.
  • Confirmed the unpatched TRT-LLM 1.3.0rc19 loader reproduces the 20480 vs. 18432 fused-QKV scale mismatch.
  • Confirmed the patched loader loads the full checkpoint, starts trtllm-serve, and successfully serves a completion request.

• Generic dense W4A8 control:

  • Exported and served TinyLlama using the same W4A8_AWQ + FP8 KV-cache recipe.
  • Confirmed the existing generic W4A8 loader and runtime path remain functional.

The B200 runs used calib_size=8 as a structural and functional smoke test; no accuracy claim is made.

Dev Engineer Review

• Updated Gemma4 HF checkpoint preprocessing to regex-match layer_scalar and k_proj keys independently of dictionary order and optional language_model. prefixes.
• For full-attention layers with attention_k_eq_v=True, missing v_proj tensors are derived from K tensors, including W4A8 k_scale/k_bias metadata mapped to V equivalents.
• Explicit V tensors remain preserved; processing is idempotent and scoped to Gemma4 checkpoint preprocessing without public API or kernel changes.
• Regression coverage includes text/VLM prefixes, full vs. sliding attention, explicit V values, scalar buffers, quantization metadata, and repeated preprocessing.
• No configuration or test-list changes were identified.

QA Engineer Review

• Added test: TestGemma4HfWeightMapper.test_duplicate_full_attention_kv_projection_tensors.
• The test covers both non-VLM and VLM mapper layouts and validates KV duplication, metadata handling, preservation of explicit V tensors, scalar copying, exclusion of sliding-attention V keys, and idempotence.
• No tests/integration/test_lists/ changes are included, so test-list coverage is not explicitly registered.
• Verdict: needs follow-up to confirm CI/manual QA coverage registration.

[coderabbitai]

Walkthrough

Changes

Gemma4 weight preprocessing

Layer / File(s)	Summary
Regex-driven buffer and KV preprocessing tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py	Regex patterns match layer scalar buffers and K-projection tensors, enabling scalar copying and conditional full-attention V-projection duplication with scale and bias suffix mapping.
Preprocessing behavior validation tests/unittest/_torch/modeling/test_modeling_gemma4.py	Parameterized tests cover VLM and non-VLM layouts, tensor preservation and omission, scalar-buffer copying, and idempotent preprocessing.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: kaiyux

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title follows the required [ticket][type] Summary pattern and clearly matches the Gemma4 W4A8 K=V checkpoint fix.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections and is detailed enough.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 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/unittest/_torch/modeling/test_modeling_gemma4.py`:
- Around line 455-457: Update the idempotency check around
mapper.preprocess_weights so it snapshots result’s keys before the second pass,
then verifies second_result contains exactly that original key set while
retaining the existing identity assertions for values.

🪄 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: 0c51b53f-6f79-4110-85ce-eb60a36b33f9

📥 Commits

Reviewing files that changed from the base of the PR and between 77774d9 and 003eb57.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
• tests/unittest/_torch/modeling/test_modeling_gemma4.py

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61327 [ run ] triggered by Bot. Commit: 003eb57 Link to invocation

[tensorrt-cicd]

PR_Github #61327 [ run ] completed with state FAILURE. Commit: 003eb57
/LLM/main/L0_MergeRequest_PR pipeline #49556 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61454 [ run ] triggered by Bot. Commit: 003eb57 Link to invocation

[tensorrt-cicd]

PR_Github #61454 [ run ] completed with state SUCCESS. Commit: 003eb57
/LLM/main/L0_MergeRequest_PR pipeline #49675 completed with status: 'SUCCESS'

CI Report

Link to invocation