[None][perf] Improve inference correctness and perf for Gemma4 - #15848
Conversation
2ez4bz
changed the title
📝 WalkthroughWalkthroughFlashInfer KV cache metadata planning shifts to host-side block counting with a new ChangesKV Cache Block Index Planning
Gemma4 Routing, Attention Masking, and Multimodal Encoding
Tokenizer DecodeStream Invalid-Prefix Recovery
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant FlashInferAttentionMetadata
participant KVCacheManager
participant KVCacheManagerV2
FlashInferAttentionMetadata->>KVCacheManager: get_batch_cache_indices(request_ids, num_blocks_per_seq)
KVCacheManager->>KVCacheManagerV2: get_batch_cache_indices(request_ids, num_blocks_per_seq)
KVCacheManagerV2-->>KVCacheManager: truncated block indices per request
KVCacheManager-->>FlashInferAttentionMetadata: paged KV block indices
sequenceDiagram
participant Forward
participant HasActiveMultimodalTokens
participant ForwardMultimodalEncoder
participant GetMultimodalEmbeddings
participant FindInputMmEmbeds
Forward->>HasActiveMultimodalTokens: filter active multimodal_params
Forward->>GetMultimodalEmbeddings: encode via ForwardMultimodalEncoder
GetMultimodalEmbeddings->>ForwardMultimodalEncoder: concat image/video/audio inputs
ForwardMultimodalEncoder-->>GetMultimodalEmbeddings: mm_embeds
GetMultimodalEmbeddings-->>Forward: mm_embeds
Forward->>FindInputMmEmbeds: align embeds with active params
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)
1323-1355: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winTruncation is applied to the packed beam layout, not the raw per-beam block list.
result[i] = result[i][:num_blocks_per_seq[i]]runs after_pack_beam_cache_indicesforbeam_width > 1. The packed layout is "beam 0's full block list + one divergent final block per other beam that differs from beam 0's", which is not a simple ordered block count — slicing it by a raw block-count value could silently drop the wrong elements (e.g. cut into beam-0's shared prefix while intending to trim trailing padding) if this method is ever called with bothbeam_width > 1andnum_blocks_per_seqset.Today only
FlashInferAttentionMetadata.prepare()calls this withnum_blocks_per_seq, and it always uses the defaultbeam_width=1, so this isn't currently exercised. Still, worth a guard/comment (or an assertion thatbeam_width == 1whennum_blocks_per_seqis provided) to prevent silent misuse later.🛡️ Suggested defensive guard
result[i] = beams[ 0] if beam_width == 1 else self._pack_beam_cache_indices(beams) if num_blocks_per_seq is not None: + assert beam_width == 1, ( + "num_blocks_per_seq truncation is only well-defined for " + "beam_width == 1 (packed beam layout is not a simple " + "block-ordered list)." + ) result[i] = result[i][:num_blocks_per_seq[i]]🤖 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 `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 1323 - 1355, The truncation in get_batch_cache_indices is applied after _pack_beam_cache_indices, which makes num_blocks_per_seq unsafe for beam_width > 1 because the packed layout is not a simple block list. Update get_batch_cache_indices to guard this path by asserting or rejecting num_blocks_per_seq unless beam_width == 1, or move truncation to the raw per-beam data before packing so the behavior stays correct and explicit.
🧹 Nitpick comments (5)
tests/unittest/_torch/modeling/test_modeling_gemma4.py (2)
790-817: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new routing-stabilization helper.
The new helper and nested hook currently infer untyped parameters/returns.
Proposed typing cleanup
- def _stabilize_moe_routing(self, hf_model, trt_model, config): + def _stabilize_moe_routing( + self, + hf_model: torch.nn.Module, + trt_model: torch.nn.Module, + config: Gemma4TextConfig, + ) -> None: @@ - def add_expert_offsets(_module, _inputs, output): + def add_expert_offsets( + _module: torch.nn.Module, + _inputs: tuple[torch.Tensor, ...], + output: torch.Tensor, + ) -> torch.Tensor: return output + expert_offsetsAs per coding guidelines, “Always annotate functions. Make the return type
Noneif the function does not return anything.”🤖 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 `@tests/unittest/_torch/modeling/test_modeling_gemma4.py` around lines 790 - 817, The new routing-stabilization helper in test_modeling_gemma4.py is missing explicit type annotations on its parameters and nested hook. Update _stabilize_moe_routing to annotate hf_model, trt_model, and config, and annotate the nested add_expert_offsets hook’s arguments and return type; also make the outer helper’s return type None since it only registers hooks and returns nothing.Source: Coding guidelines
2028-2080: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the new tests and local stub.
Coverage looks sufficient in
tests/unittest/_torch/modeling/test_modeling_gemma4.pyfor this routing/mask layer; this is just to keep new test code typed.Proposed typing cleanup
- def test_bidirectional_mask_only_applies_to_sliding_layers(self): + def test_bidirectional_mask_only_applies_to_sliding_layers(self) -> None: @@ - def passthrough(*, hidden_states, **_kwargs): + def passthrough(*, hidden_states: torch.Tensor, **_kwargs: object) -> torch.Tensor: return hidden_states @@ - def test_gemma4_routing_matches_hf_reference(self): + def test_gemma4_routing_matches_hf_reference(self) -> None:As per coding guidelines, “Always annotate functions.” As per path instructions, review
tests/**changes for coverage sufficiency.🤖 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 `@tests/unittest/_torch/modeling/test_modeling_gemma4.py` around lines 2028 - 2080, The new test methods and the local passthrough stub in test_modeling_gemma4 need explicit type annotations to satisfy the “Always annotate functions” guideline. Add return and parameter annotations for the test functions in Gemma4-related test cases, and annotate the nested passthrough helper so its signature is typed while still matching the layer.forward call shape. Keep the changes confined to the test helpers and methods around test_bidirectional_mask_only_applies_to_sliding_layers and test_gemma4_routing_matches_hf_reference.Sources: Coding guidelines, Path instructions
tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py (1)
400-422: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMutable dict class attribute flagged by static analysis.
EXTRA_EVALUATOR_KWARGS = {...}is a mutable default at class scope, flagged by Ruff (RUF012). Note: per a retrieved learning, this repo's Ruffselectis limited to{D, E, F, I, PLE, W}, which does not includeRUF-prefixed rules — so this may not actually be enforced by the lint gate here. Flagging for awareness only.♻️ Optional fix
- EXTRA_EVALUATOR_KWARGS = { - "chat_template_kwargs": {"enable_thinking": False}, - } + EXTRA_EVALUATOR_KWARGS: ClassVar[dict] = { + "chat_template_kwargs": {"enable_thinking": False}, + }Based on learnings, "Ruff is configured in pyproject.toml to enable only {D, E, F, I, PLE, W}", so this specific rule category may not be enforced repo-wide.
🤖 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 `@tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py` around lines 400 - 422, The class-level EXTRA_EVALUATOR_KWARGS on TestGemma4_26B_A4B is a mutable dict that Ruff’s RUF012 would flag; make it immutable at class scope by replacing it with a non-mutable pattern or constructing it per instance in the test harness path, keeping the same chat_template_kwargs content. Update the TestGemma4_26B_A4B definition only, and preserve the existing sampling_params and kv_cache_config behavior.Sources: Learnings, Linters/SAST tools
tests/unittest/_torch/modeling/test_gemma4_multimodal.py (1)
705-758: 🎯 Functional Correctness | 🔵 TrivialGood coverage for single-modality chunking; consider adding a mixed-modality ordering test.
test_chunked_prefill_reuses_cached_vision_embeddingsandtest_chunk_without_multimodal_tokens_is_inactivecorrectly validate the chunk-slicing math and inactive-chunk short-circuit for a single image-only param.Coverage is currently insufficient for the scenario where a batch has multiple active
multimodal_paramswith different modality types (e.g. one audio-only param followed by one image-only param). Per the concern raised inmodeling_gemma4mm.py's_forward_multimodal_encoder(lines 902-994), such a batch could produce embeddings in a different order than the token sequence expects. Recommend adding a test totests/unittest/_torch/modeling/test_gemma4_multimodal.pythat exercisesget_multimodal_embeddings/_forward_multimodal_encoderwith two params of different modalities and asserts embedding-to-token alignment.🤖 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 `@tests/unittest/_torch/modeling/test_gemma4_multimodal.py` around lines 705 - 758, Add coverage for mixed-modality batching in the multimodal encoder path, since the current tests only validate single-image chunking and inactive chunks. Create a test in test_gemma4_multimodal.py that drives get_multimodal_embeddings through _forward_multimodal_encoder with two active multimodal_params of different types (for example audio-only then image-only) and assert the returned embeddings preserve token-sequence order. Use the existing helpers and symbols like MultimodalParams, MultimodalRuntimeData, get_multimodal_embeddings, and find_input_mm_embeds to verify alignment.Source: Path instructions
tensorrt_llm/_torch/attention_backend/flashinfer.py (1)
911-938: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant re-truncation after
get_batch_cache_indicesalready truncates.
get_batch_cache_indices(..., num_blocks_per_seq=self.num_blocks)now already truncates each per-request block list server-side, soblock_ids[:self.num_blocks[i]]here (and the analogous slice in the VSWA branch below at line 967) is redundant — harmless since slicing beyond length is a no-op, but slightly misleading since it implies the returned lists might still be longer thannum_blocks[i].♻️ Optional cleanup
- paged_kv_indices_list = [] - for i, block_ids in enumerate(block_ids_per_seq): - paged_kv_indices_list.extend(block_ids[:self.num_blocks[i]]) + paged_kv_indices_list = [ + idx for block_ids in block_ids_per_seq for idx in block_ids + ]🤖 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 `@tensorrt_llm/_torch/attention_backend/flashinfer.py` around lines 911 - 938, `get_batch_cache_indices` already truncates each per-request block list, so the extra slicing in the `paged_kv_indices_list` build inside the cache update path is redundant and misleading. Remove the `[:self.num_blocks[i]]` truncation in this loop, and make the same cleanup in the analogous VSWA branch so the `block_ids_per_seq` handling matches the actual contract of `get_batch_cache_indices`.
🤖 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 `@tensorrt_llm/_torch/models/modeling_gemma4mm.py`:
- Around line 902-1019: The multimodal embeddings produced by
_forward_multimodal_encoder are currently concatenated by modality type instead
of the original multimodal_params/request order, which can misalign cached
slices later. Update _forward_multimodal_encoder and the downstream cache/slice
flow (including _cache_multimodal_embeddings and any use of
get_multimodal_embeddings/find_input_mm_embeds) so embeddings are emitted or
reassembled in per-request order, preserving each multimodal_param’s row mapping
across mixed image/video/audio batches.
In `@tests/unittest/_torch/executor/test_per_layer_head_dim.py`:
- Around line 22-26: The current coverage only exercises KVCacheManagerV2, so
add a matching V1 test in test_resource_manager.py that verifies
num_blocks_per_seq after beam packing with beam_width > 1 and truncation
behavior. Use the existing KVCacheManager test setup/helpers in that file to
build the case, then assert the expected num_blocks_per_seq values after
packing/truncation so the V1 path is covered alongside the new V2 test.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 1323-1355: The truncation in get_batch_cache_indices is applied
after _pack_beam_cache_indices, which makes num_blocks_per_seq unsafe for
beam_width > 1 because the packed layout is not a simple block list. Update
get_batch_cache_indices to guard this path by asserting or rejecting
num_blocks_per_seq unless beam_width == 1, or move truncation to the raw
per-beam data before packing so the behavior stays correct and explicit.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/flashinfer.py`:
- Around line 911-938: `get_batch_cache_indices` already truncates each
per-request block list, so the extra slicing in the `paged_kv_indices_list`
build inside the cache update path is redundant and misleading. Remove the
`[:self.num_blocks[i]]` truncation in this loop, and make the same cleanup in
the analogous VSWA branch so the `block_ids_per_seq` handling matches the actual
contract of `get_batch_cache_indices`.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py`:
- Around line 400-422: The class-level EXTRA_EVALUATOR_KWARGS on
TestGemma4_26B_A4B is a mutable dict that Ruff’s RUF012 would flag; make it
immutable at class scope by replacing it with a non-mutable pattern or
constructing it per instance in the test harness path, keeping the same
chat_template_kwargs content. Update the TestGemma4_26B_A4B definition only, and
preserve the existing sampling_params and kv_cache_config behavior.
In `@tests/unittest/_torch/modeling/test_gemma4_multimodal.py`:
- Around line 705-758: Add coverage for mixed-modality batching in the
multimodal encoder path, since the current tests only validate single-image
chunking and inactive chunks. Create a test in test_gemma4_multimodal.py that
drives get_multimodal_embeddings through _forward_multimodal_encoder with two
active multimodal_params of different types (for example audio-only then
image-only) and assert the returned embeddings preserve token-sequence order.
Use the existing helpers and symbols like MultimodalParams,
MultimodalRuntimeData, get_multimodal_embeddings, and find_input_mm_embeds to
verify alignment.
In `@tests/unittest/_torch/modeling/test_modeling_gemma4.py`:
- Around line 790-817: The new routing-stabilization helper in
test_modeling_gemma4.py is missing explicit type annotations on its parameters
and nested hook. Update _stabilize_moe_routing to annotate hf_model, trt_model,
and config, and annotate the nested add_expert_offsets hook’s arguments and
return type; also make the outer helper’s return type None since it only
registers hooks and returns nothing.
- Around line 2028-2080: The new test methods and the local passthrough stub in
test_modeling_gemma4 need explicit type annotations to satisfy the “Always
annotate functions” guideline. Add return and parameter annotations for the test
functions in Gemma4-related test cases, and annotate the nested passthrough
helper so its signature is typed while still matching the layer.forward call
shape. Keep the changes confined to the test helpers and methods around
test_bidirectional_mask_only_applies_to_sliding_layers and
test_gemma4_routing_matches_hf_reference.
🪄 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: 3aa06af3-1969-43fa-ad0a-2a55ec2633eb
📒 Files selected for processing (13)
tensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/models/modeling_gemma4.pytensorrt_llm/_torch/models/modeling_gemma4mm.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/tokenizer/tokenizer.pytests/integration/defs/accuracy/references/mmmu.yamltests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/executor/test_per_layer_head_dim.pytests/unittest/_torch/modeling/test_gemma4_multimodal.pytests/unittest/_torch/modeling/test_modeling_gemma4.pytests/unittest/llmapi/test_tokenizer_decode.py
|
[by Codex] @yizhang-nv Friendly reminder: could you please review this PR? Thanks! |
* Why? Gemma 4 diverged from the Hugging Face reference for multimodal attention and MoE routing, while chunked prefill repeated avoidable encoder and mask work. Incremental decoding could also fail on invalid byte prefixes, and padded KV page tables added unnecessary overhead. * What? This commit: - aligns Gemma 4 masks and routing with Hugging Face, streamlines mask construction, and reuses multimodal embeddings across prefill chunks. - recovers invalid HF decode streams and limit FlashInfer and KV cache page processing to each sequence's live block count. - tweaks unit tests against the HF reference implementation to tighten bounds by an order of magnitude. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
|
/bot run |
|
/bot run --disable-fail-fast |
|
PR_Github #57304 [ run ] triggered by Bot. Commit: |
|
PR_Github #57304 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
|
/bot run |
|
PR_Github #57637 [ run ] triggered by Bot. Commit: |
|
PR_Github #57637 [ run ] completed with state |
…nfer plans, CUDA graph padding prealloc Combine three groups of Gemma4 single-GPU serving optimizations on top of the page-table work from NVIDIA#15848: - FlashInfer host-path: build all page-table metadata (kv lens, block counts, indptrs, last-page lens) with numpy on the host and stage through pinned tensors, add a flat KV cache indices path (get_batch_cache_indices_flat) that never materializes padded per-request lists, retain host copies of decode indptr and pool indices, and pre-build persistent trtllm-gen decode block tables on the host so decode plans need no GPU round trips; vectorize the VSWA decode block-table refresh. - Fused Gemma4 modules: model-side fused QKV prep (norm + RoPE + FP8 quant), fused GeLU-mul + NVFP4 quant, fused RMSNorm + NVFP4 quant, and a fused residual-add tail, wired into modeling_gemma4 with per-fusion opt-out env vars (TRTLLM_GEMMA4_DISABLE_*); FlashInfer emits BF16 output when Q arrives pre-quantized to FP8. - CUDA graph padding: pre-allocate the padding dummy request at warmup while the KV cache still has free blocks, so padded batches cannot permanently fall back to eager mode once the cache saturates. Benchmark (1xB200, Gemma-4-31B-IT-NVFP4, ISL 1024 / OSL 128, 1024 prompts, concurrency 1024): output throughput 1324 -> 1733 tok/s (+30.9%), median TTFT -23.6%, median TPOT -24.4%, median E2EL -23.9%. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…A#15848) * Why? Gemma 4 diverged from the Hugging Face reference for multimodal attention and MoE routing, while chunked prefill repeated avoidable encoder and mask work. Incremental decoding could also fail on invalid byte prefixes, and padded KV page tables added unnecessary overhead. * What? This commit: - aligns Gemma 4 masks and routing with Hugging Face, streamlines mask construction, and reuses multimodal embeddings across prefill chunks. - recovers invalid HF decode streams and limit FlashInfer and KV cache page processing to each sequence's live block count. - tweaks unit tests against the HF reference implementation to tighten bounds by an order of magnitude. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
9 participants
Summary by CodeRabbit
Description
Gemma 4 diverged from the Hugging Face reference for multimodal
attention and MoE routing, while chunked prefill repeated avoidable
encoder and mask work. Incremental decoding could also fail on invalid
byte prefixes, and padded KV page tables added unnecessary overhead.
This commit:
construction, and reuses multimodal embeddings across prefill chunks.
page processing to each sequence's live block count.
bounds by an order of magnitude.
Performance comparison
NVFP4, B200
26B
31B
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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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
To see a list of available CI bot commands, please comment
/bot help.