[https://nvbugs/6185146][fix] Use mat_a.new_empty([m, n_out//2]) / input_scale.new_empty([sf_size]) in the - #14710
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 (3)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthroughThis PR optimizes the FP4-output CuTe DSL kernel path for torch.compile by updating memory allocation in the fake custom op and adjusting kernel selection logic in GatedMLP to avoid SymBool guards, enabling correct CUDA graph piecewise capture. A passing test waiver is removed. ChangesFP4 Output Kernel Path Optimization
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
3cc1518 to
70e3072
Compare
|
Fix the DCO issue with sign-off |
70e3072 to
2e39ffe
Compare
|
/bot run |
|
PR_Github #53022 [ run ] triggered by Bot. Commit: |
|
PR_Github #53022 [ run ] completed with state
|
2e39ffe to
9714de2
Compare
|
/bot run |
|
PR_Github #54444 [ run ] triggered by Bot. Commit: |
|
PR_Github #54444 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54494 [ run ] triggered by Bot. Commit: |
|
PR_Github #54494 [ run ] completed with state
|
9714de2 to
b418db4
Compare
|
/bot run |
|
PR_Github #55117 [ run ] triggered by Bot. Commit: |
|
PR_Github #55117 [ run ] completed with state
|
…d fp4out shared-expert path The CuTe-DSL fused GEMM+SwiGLU+quant shared-expert path (PR NVIDIA#11897) broke TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus with torch_compile=True + piecewise CUDA graph. piecewise_optimizer raised "Cannot identify dynamic shape" because, on the m >= 128 branch: 1. The fake function for cute_dsl_nvfp4_dense_gemm_swiglu_fp4out_blackwell used positional torch.empty(m, n_out // 2, ...), which did not propagate the input's token-dim SymInt into the FX graph the same way that mat_a.new_empty([...]) does (matching the BF16 fake and nvfp4_gemm fake). 2. fp4_out=m >= GatedMLP._FP4OUT_MIN_M produced a SymBool that, when consumed by `if fp4_out:` in _fused_gate_up_swiglu, forced a guard on the dynamic token dim during torch.compile + piecewise capture. Fix: switch the fake function to mat_a.new_empty / input_scale.new_empty, and skip the data-dependent m >= MIN_M check while compiling (the kernel already pads m up to the CTA tile height internally, so fp4_out is safe at any m). Also remove the corresponding waives.txt entry. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com>
Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com>
b418db4 to
9c32fb4
Compare
|
/bot run |
|
PR_Github #55171 [ run ] triggered by Bot. Commit: |
|
PR_Github #55171 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55213 [ run ] triggered by Bot. Commit: |
|
PR_Github #55213 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55232 [ run ] triggered by Bot. Commit: |
|
PR_Github #55232 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55367 [ run ] triggered by Bot. Commit: |
|
PR_Github #55367 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55392 [ run ] triggered by Bot. Commit: |
|
PR_Github #55392 [ run ] completed with state |
…`input_scale.new_empty([sf_size])` in the (NVIDIA#14710) Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
…`input_scale.new_empty([sf_size])` in the (NVIDIA#14710) Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
4 participants
Summary
m >= 128(creating a SymBool guard), so piecewise_optimizer could not find a placeholder SymInt equal to compile_time_num_tokens for the m>=128 compile branch.mat_a.new_empty([m, n_out//2])/input_scale.new_empty([sf_size])in the fp4out fake fn, and skipm >= _FP4OUT_MIN_Mundertorch.compiler.is_compiling()(kernel pads m to CTA tile internally, so fp4_out=True is safe). Remove the now-passing waiver.Test plan
Links
Summary by CodeRabbit
Release Notes
Refactor
Tests