[https://nvbugs/6737351][fix] Remove CGA smem-reduction clamp from trtllm-gen FMHA kernel selection - #18945
[https://nvbugs/6737351][fix] Remove CGA smem-reduction clamp from trtllm-gen FMHA kernel selection#18945kris1025 wants to merge 1 commit into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughFMHA validation, warmup, and execution now call ChangesFMHA kernel selection
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to FMHA validation, warmup, and execution now select kernels directly through the autotuner, restoring unconstrained CTA selection. No actionable merge-blocking risk is currently identified. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
kris1025
changed the title
…el selection PR NVIDIA#18518 wrapped every FmhaAutoTuner::selectKernel() call in selectKernelWithCgaSmemReductionLimit(), which clamps FmhaOptions::mMaxNumCtasPerSeqKv from INT_MAX to 8 for any GmemReduction + Static-scheduler request on ALL architectures. That request pattern is exactly the trtllm-gen generation (decode) path (xqaDispatcher sets multi_block_mode -> GmemReduction + Static). On B200 (SM100) MLA decode this changes the tuned kernel selection: - computeNumCtas caps the split-KV CTA count at 8 instead of filling the machine in one wave (e.g. ~18 splits for 8-9k KV at batch 1); - the smaller grid flips selectMlaGenerationKernel's headDimPerCtaV fine-tuning from 256 to 128, duplicating the 576-dim QK^T work 4x instead of 2x; - the one-wave check in enableCgaReduction now succeeds and promotes GmemReduction to CgaSmemReduction with far fewer CTAs. Net effect: decode attention loses most of its split-KV parallelism at low batch, regressing disagg gen-only DeepSeek-R1 FP4 MTP3 at con1 (disagg-gen_only-b200_deepseek-r1-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_ eplb0_mtp3_ccb-NIXL). Remove the wrapper and restore the direct selectKernel() calls, returning SM100-family kernel selection to the pre-PR behavior. The clamp existed to keep a Rubin promotion path legal; that constraint belongs in the trtllm-gen autotuner (or an SM107-gated check), not in arch-agnostic host code. Signed-off-by: linquanh <linquanh@nvidia.com>
5287633 to
5ed497f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #72371 [ run ] triggered by Bot. Commit: |
|
/bot run --stage-list "DGX_B200-16_GPUs-2_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE1-GPU8-Post-Merge" |
|
PR_Github #72376 [ run ] triggered by Bot. Commit: |
|
PR_Github #72371 [ run ] completed with state |
|
PR_Github #72376 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #72397 [ run ] triggered by Bot. Commit: |
|
PR_Github #72397 [ run ] completed with state
|
2 participants
Description
Problem: PR #18518 wrapped every
FmhaAutoTuner::selectKernel()call inselectKernelWithCgaSmemReductionLimit(), which clampsFmhaOptions::mMaxNumCtasPerSeqKvfromINT_MAXto16 / max(clusterDimX, 2)(= 8) for anyGmemReduction+Static-scheduler request — on all architectures. That request pattern is exactly the trtllm-gen generation (decode) path:xqaDispatchersetsmMultiCtasKvMode = multi_block_mode(→GmemReduction) andmTileScheduler = Staticwhen multi-block mode is on.On B200 (SM100) MLA decode this silently changes the tuned kernel selection:
computeNumCtascaps the split-KV CTA count at 8 instead of filling the machine in one wave (numCtasPerSeqKv = min(maxNumCtasPerSeqKv, SM_count / (numCtasX*numCtasY*numCtasZ)); e.g. ~18 splits for 8–9k KV at batch 1 on 148 SMs).selectMlaGenerationKernel'sheadDimPerCtaVfine-tuning from 256 to 128 (totalNumCtas * 4 <= multiProcessorCountnow passes), so the 576-dim QK^T work is duplicated 4× (V-split) instead of 2×.enableCgaReductionnow succeeds and promotesGmemReduction→CgaSmemReductionwith far fewer CTAs.Net effect: decode attention loses most of its split-KV parallelism at low batch. This regresses disagg gen-only DeepSeek-R1 FP4 MTP3 at con1:
disagg-gen_only-b200_deepseek-r1-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL [con1_iter1_isl8192_osl1024].High-concurrency cases are unaffected because the batch×heads grid dimension already fills the GPU; the ctx/prefill path is unaffected because it uses the persistent scheduler with
multiCtasKvModedisabled.Solution: remove the wrapper (
selectKernelWithCgaSmemReductionLimit/limitCgaSmemReductionCtasKv) and restore the directFmhaAutoTuner::selectKernel()calls at all three sites (checkIfKernelExist,warmupOneKernel,run), returning SM100-family kernel selection to the exact pre-#18518 behavior.The clamp was introduced to keep a Rubin (SM107) CGA-promotion path legal (
computeNumCtaschecksnumCtasPerSeqKv * clusterDimX <= 16). That constraint belongs in the trtllm-gen autotuner itself (or behind anmSM == kSM_107gate), not in arch-agnostic host code that overrides tuned SM100 heuristics. Note SM107 is not reachable in current public builds (TLLM_RUBIN_FEATURESis not defined and Sm107a cubins are not archived), so this removal does not change any shipping Rubin behavior.Test Coverage
tests/unittest/_torch/attention/test_combined_fmha.py,test_fmha_page_index.py.perf/test_perf_sanity.py::test_e2e[disagg-gen_only-b200_deepseek-r1-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL](the case that regressed) should return to its pre-[None][feat] Enable Rubin (SM107) trtllm-gen FMHA features #18518 baseline.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why.
PR follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (N/A — pure revert of host-side logic to the pre-regression state; covered by existing tests above).
No API changes, no new dependencies, no CODEOWNERS changes, no documentation impact.
Please check this after reviewing the above items as appropriate for this PR.
Dev Engineer Review
FmhaAutoTuner::selectKernel()calls for validation, JIT warmup, and execution.mMaxNumCtasPerSeqKvto 8.CgaSmemReduction.QA Engineer Review
No test changes.
Per-File QA Perspective
cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h: Kernel selection no longer applies the shared-memory reduction clamp or post-selection CTA limit. Verify FMHA behavior, kernel tuning, and low-batch B200 MLA decode performance across affected architectures.