[None][perf] Optimize Gemma4 serving: fused kernels, sync-free FlashInfer plans, CUDA graph padding prealloc - #16033
Closed
kaiyux wants to merge 3 commits into
Closed
Conversation
…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>
Contributor
📝 WalkthroughWalkthroughThis PR introduces sync-free host-side decode planning for the FlashInfer attention backend (pinned staging buffers, vectorized block-table construction, CUDA-graph updates), a CUDA-graph padding-dummy preallocation mechanism, a ChangesFlashInfer Sync-Free Decode Planning & CUDA Graph Padding
Gemma4 Fused Triton Kernels
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant ModelEngine as PyTorchModelEngine
participant CudaGraphRunner
participant KVCacheManager
ModelEngine->>CudaGraphRunner: preallocate_padding_dummy(resource_manager)
CudaGraphRunner->>KVCacheManager: allocate dummy request (draft_len=0)
KVCacheManager-->>CudaGraphRunner: dummy request or failure
CudaGraphRunner-->>ModelEngine: cached in padding_dummy_requests
Note over ModelEngine,CudaGraphRunner: Later, during a padded step
ModelEngine->>CudaGraphRunner: _get_padded_batch(runtime_draft_len)
CudaGraphRunner->>CudaGraphRunner: _get_or_create_padding_dummy(runtime_draft_len)
alt cached dummy exists
CudaGraphRunner-->>ModelEngine: reuse cached dummy request
else allocation fails
CudaGraphRunner-->>ModelEngine: warning_once + eager fallback (0)
end
sequenceDiagram
participant Attention as Gemma4Attention
participant FusedQKV as gemma4_fused_qkv_norm_rope_quant
participant DecoderLayer as Gemma4DecoderLayer
participant FusedTail as gemma4_fused_norm_add_scale
participant MLP as _Gemma4GeluQuantMLP
Attention->>Attention: _fused_qkv_prep_enabled() check
Attention->>FusedQKV: apply_rope(qkv, position_ids, cos_sin)
FusedQKV-->>Attention: fp8 Q, K, V
DecoderLayer->>DecoderLayer: _fused_norm_quant_enabled() check
DecoderLayer->>MLP: Fp4QuantizedTensor input
MLP-->>DecoderLayer: fp4-quantized MLP output
DecoderLayer->>FusedTail: gemma4_fused_norm_add_scale(hidden, residual, next_norm_weight)
FusedTail-->>DecoderLayer: (hidden_states, pre_normed)
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
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}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (4)
tensorrt_llm/_torch/models/gemma4_fused_qkv.py (1)
175-175: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider asserting the
head_dimpower-of-two assumption.
block_n = max(1, min(32, 4096 // head_dim))is only guaranteed power-of-two (required bytl.arangeandHALF) whenhead_dimis a power-of-two divisor of 4096. That holds for the gated shapes today, but a strayhead_dim(e.g. 192) would silently produce a non-power-of-twoblock_n/halfand a confusing kernel-compile failure. A cheapassert (head_dim & (head_dim - 1)) == 0would fail fast with a clear message.🤖 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/models/gemma4_fused_qkv.py` at line 175, The block_n calculation in gemma4_fused_qkv relies on head_dim being a power-of-two divisor of 4096, but this assumption is not checked and can lead to invalid kernel settings and opaque compile failures. Add a fast-fail assertion near the block_n setup in the relevant fused QKV path to verify head_dim is power-of-two before computing block_n, and include a clear error message so the failure is explicit and easy to trace.tensorrt_llm/_torch/models/gemma4_fused_gelu_quant.py (1)
42-42: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer built-in
tupleovertyping.Tuple.The codebase targets Python 3.10+, so annotate with
tuple[...]and drop thetyping.Tupleimport (applies to the sibling fused-kernel modulesgemma4_fused_norm_quant.py,gemma4_fused_qkv.py, andgemma4_fused_tail.pyas well).As per coding guidelines: "Prefer built-in types
list,dict,tupleovertyping.List,typing.Dict,typing.Tuple".🤖 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/models/gemma4_fused_gelu_quant.py` at line 42, The module still imports and uses typing.Tuple even though the codebase targets Python 3.10+, so update the type annotations in gemma4_fused_gelu_quant and the sibling fused-kernel modules gemma4_fused_norm_quant, gemma4_fused_qkv, and gemma4_fused_tail to use built-in tuple[...] syntax instead. Remove the Tuple import from each module after switching the relevant return/type annotations so the code follows the project guideline to prefer built-in collection types.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/flashinfer.py (1)
922-948: 🩺 Stability & Availability | 🔵 Trivial | ⚖️ Poor tradeoffPersistent per-
plan_paramsbuffers are never reclaimed.
_decode_block_tables_dev/_decode_block_tables_stageaccumulate a(max_num_requests, width)int32 device buffer (plus pinned stage) per distinct non-maskedPlanParams._clean_cached_plansdrops masked entries from_plan_params_to_wrappersbut nothing prunes these two dicts, so any churn in plan params (e.g. varyingattention_window_size/dtype across steps) slowly leaks device memory. In steady state plan params are few, so impact is limited; consider pruning entries whoseplan_paramsno longer appear in_plan_params_to_wrappers.🤖 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 922 - 948, Persistent per-PlanParams decode buffers are never reclaimed, so _decode_block_tables_dev and _decode_block_tables_stage can grow indefinitely as plan params change. Update the cache management in FlashinferBackend to prune these entries when _clean_cached_plans removes stale plans, using the existing _plan_params_to_wrappers / plan_params keys to identify buffers no longer in use. Keep the cleanup aligned with the decode buffer allocation path in the plan/replan logic so both the device tensor and pinned stage buffer are dropped together.tensorrt_llm/_torch/models/modeling_gemma4.py (1)
824-838: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDe-duplicate the FP4 pre-quant eligibility check.
The five "pre-quantized-input path" checks here on
gate_up_proj(Lines 833-837) are identical to those ondown_projin_Gemma4GeluQuantMLP._fused_gelu_quant_enabled(Lines 137-141). Since both intentionally mirrorLinear._input_prepare, extracting a small module-level helper keeps the two call sites from silently diverging if that contract changes.♻️ Suggested helper
def _linear_accepts_prequant_fp4(linear) -> bool: # Mirror the checks the pre-quantized-input path in # Linear._input_prepare enforces for Fp4QuantizedTensor. return ( getattr(linear, "has_nvfp4", False) and not getattr(linear, "force_dynamic_quantization", True) and getattr(linear, "input_scale", None) is not None and getattr(linear, "pre_quant_scale", None) is None and getattr(linear, "scaling_vector_size", None) == 16 )Then both predicates reduce to
... and _linear_accepts_prequant_fp4(dp)/_linear_accepts_prequant_fp4(gu).🤖 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/models/modeling_gemma4.py` around lines 824 - 838, De-duplicate the repeated FP4 pre-quant eligibility logic used in `_fused_norm_quant_enabled` and `_Gemma4GeluQuantMLP._fused_gelu_quant_enabled` by introducing a small shared helper for the `Linear._input_prepare`-style checks. Move the common `gate_up_proj`/`down_proj` predicate into a module-level function (for example, one that validates `has_nvfp4`, `force_dynamic_quantization`, `input_scale`, `pre_quant_scale`, and `scaling_vector_size`) and call it from both `gu` and `dp` sites so the two paths stay in sync.
🤖 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 `@tensorrt_llm/_torch/attention_backend/flashinfer.py`:
- Around line 922-948: Persistent per-PlanParams decode buffers are never
reclaimed, so _decode_block_tables_dev and _decode_block_tables_stage can grow
indefinitely as plan params change. Update the cache management in
FlashinferBackend to prune these entries when _clean_cached_plans removes stale
plans, using the existing _plan_params_to_wrappers / plan_params keys to
identify buffers no longer in use. Keep the cleanup aligned with the decode
buffer allocation path in the plan/replan logic so both the device tensor and
pinned stage buffer are dropped together.
In `@tensorrt_llm/_torch/models/gemma4_fused_gelu_quant.py`:
- Line 42: The module still imports and uses typing.Tuple even though the
codebase targets Python 3.10+, so update the type annotations in
gemma4_fused_gelu_quant and the sibling fused-kernel modules
gemma4_fused_norm_quant, gemma4_fused_qkv, and gemma4_fused_tail to use built-in
tuple[...] syntax instead. Remove the Tuple import from each module after
switching the relevant return/type annotations so the code follows the project
guideline to prefer built-in collection types.
In `@tensorrt_llm/_torch/models/gemma4_fused_qkv.py`:
- Line 175: The block_n calculation in gemma4_fused_qkv relies on head_dim being
a power-of-two divisor of 4096, but this assumption is not checked and can lead
to invalid kernel settings and opaque compile failures. Add a fast-fail
assertion near the block_n setup in the relevant fused QKV path to verify
head_dim is power-of-two before computing block_n, and include a clear error
message so the failure is explicit and easy to trace.
In `@tensorrt_llm/_torch/models/modeling_gemma4.py`:
- Around line 824-838: De-duplicate the repeated FP4 pre-quant eligibility logic
used in `_fused_norm_quant_enabled` and
`_Gemma4GeluQuantMLP._fused_gelu_quant_enabled` by introducing a small shared
helper for the `Linear._input_prepare`-style checks. Move the common
`gate_up_proj`/`down_proj` predicate into a module-level function (for example,
one that validates `has_nvfp4`, `force_dynamic_quantization`, `input_scale`,
`pre_quant_scale`, and `scaling_vector_size`) and call it from both `gu` and
`dp` sites so the two paths stay in sync.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4f345f2f-912e-4dc1-8564-ba9446b121ee
📒 Files selected for processing (14)
tensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/models/gemma4_fused_gelu_quant.pytensorrt_llm/_torch/models/gemma4_fused_norm_quant.pytensorrt_llm/_torch/models/gemma4_fused_qkv.pytensorrt_llm/_torch/models/gemma4_fused_tail.pytensorrt_llm/_torch/models/modeling_gemma4.pytensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/_torch/modules/test_gemma4_fused_gelu_quant.pytests/unittest/_torch/modules/test_gemma4_fused_norm_quant.pytests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.pytests/unittest/_torch/modules/test_gemma4_fused_tail.py
_torch/models/ holds only modeling_*.py files; architecture-specific building blocks belong under _torch/modules/ (like mhc/, engram/, fused_moe/). Move the four gemma4_fused_* kernel files into a new modules/gemma4/ subpackage and update imports in modeling_gemma4 and the unit tests. No functional change. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
The fused paths (gelu+quant, qkv prep, norm+add, norm+quant, layer tail, tail norm2) are now unconditional wherever structurally applicable; the unfused code remains only for configurations the kernels do not support (MoE block, PLE, KV-shared layers, non-NVFP4/FP8 configs, LoRA, torch.compile, custom-mask prefill). Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
kaiyux
force-pushed
the
perf-optimize/gemma-4-combined-20260706
branch
from
July 7, 2026 07:05
47ad43c to
42ea088
Compare
kaiyux
marked this pull request as ready for review
July 7, 2026 15:35
Member
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Description
Combines three groups of Gemma4 single-GPU serving optimizations, building on the page-table trimming introduced in #15848. Measured together on 1x B200 (Gemma-4-31B-IT-NVFP4, ISL 1024 / OSL 128, 1024 prompts, concurrency 1024,
trtllm-serve+benchmark_serving):1. FlashInfer host-path preparation (
flashinfer.py,kv_cache_manager_v2.py)prepare().KVCacheManagerV2.get_batch_cache_indices_flat(): gathers only the live block-table widths straight into a pinned CPU tensor with one vectorized transform, instead of materializing padded per-request Python lists.BatchDecodeWithPagedKVCacheWrapper.plan()no longer triggers flashinfer's per-request rebuild loop (one stream sync + one scalar D2H per generation request per plan). Buffers are allocated once at capacity under CUDA graphs so captured decode kernels keep reading stable addresses.2. Fused Gemma4 modules (
modeling_gemma4.py+ newgemma4_fused_*.py)TRTLLM_GEMMA4_DISABLE_FUSED_QKV_PREP,..._FUSED_GELU_QUANT,..._FUSED_NORM_QUANT,..._FUSED_TAIL,..._FUSED_TAIL_NORM2,..._FUSED_NORM_ADD).QkvE4m3OBfloat16cubins emit BF16).3. CUDA graph padding-dummy preallocation (
cuda_graph_runner.py,model_engine.py)_get_or_create_padding_dummy()(keeping the encoder-decoder handling) and logs awarning_oncewhen padding falls back to eager.Test Coverage
tests/unittest/_torch/modules/test_gemma4_fused_qkv_prep.py,test_gemma4_fused_gelu_quant.py,test_gemma4_fused_norm_quant.py,test_gemma4_fused_tail.py— parity checks against the unfused reference paths.tests/unittest/_torch/executor/test_pytorch_model_engine.pyupdated to cover the warmup padding-dummy preallocation (asserts the dummy exists after warmup and that no blocks leak beyond it)._torchattention and executor test suites;get_batch_cache_indiceskeeps its trimmed-list semantics for non-flat callers.PR Checklist
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.🤖 Generated with Claude Code