Fix architecture-independent SDPA repro failures - #386
Conversation
hwanseoc
changed the title
|
@CodeRabbit review |
✅ Action performedReview finished.
|
📝 WalkthroughWalkthroughThe PR expands SDPA tensor dumping and adds performance-mode control to FP8 and MXFP8 tests, including optional reference computation, compact MXFP8 quantization, updated validation, and a dedicated performance test. ChangesTensor dump execution
FP8 performance paths
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant exec_sdpa_mxfp8
participant quantize_to_mxfp8
participant cuDNN
participant compute_ref
exec_sdpa_mxfp8->>quantize_to_mxfp8: quantize with_ref=not perf
quantize_to_mxfp8->>cuDNN: provide MXFP8 tensors and swizzled scales
exec_sdpa_mxfp8->>cuDNN: execute forward and backward graphs
alt perf disabled
exec_sdpa_mxfp8->>compute_ref: compute reference outputs and gradients
compute_ref-->>exec_sdpa_mxfp8: return references and amax values
end
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: 1
🧹 Nitpick comments (1)
test/python/sdpa/mxfp8.py (1)
35-57: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winChunked MXFP8 quantization relies on TE's private tensor attributes.
quantize_mxfp8_compactreaches intoresult._rowwise_data,result._rowwise_scale_inv,result._columnwise_data,result._columnwise_scale_inv— underscore-prefixed (private/internal) attributes of the TE-returned tensor — to manually reconstruct anMXFP8Tensor. TheMXFP8_QUANTIZER_MAX_ROWS = 65535 * 64constant does line up with CUDA's 65535 grid-y/z launch limit, and the reconstructed constructor kwargs (rowwise_data,rowwise_scale_inv, etc.) match TE'sMXFP8TensorStorage/MXFP8Tensorpublic constructor signature, so the approach is plausible. However, depending on private instance attributes of a quantizer's return value is fragile against internal TE changes across versions (chunk-by-chunk quantization + concatenation of scale/data isn't part of TE's public API).Consider filing/tracking this as a known coupling point, or checking if TE exposes any higher-level API for chunked quantization to reduce future breakage risk when TE internals change.
🤖 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 `@test/python/sdpa/mxfp8.py` around lines 35 - 57, Review quantize_mxfp8_compact for a supported Transformer Engine API for chunked MXFP8 quantization and concatenation, replacing direct access to the result private attributes where possible. If no public API exists, document or track this coupling explicitly while preserving the current MXFP8Tensor reconstruction behavior.
🤖 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 `@include/cudnn_frontend/node/scaled_dot_product_flash_attention.h`:
- Around line 2188-2192: Update collect_tensors_to_dump_subtree() so tensor UIDs
are deduplicated across the entire subtree, not only within each node’s local
collection. Thread a shared std::unordered_set<Tensor_attributes::uid_t> through
recursive collection or deduplicate tensors_to_dump before final serialization,
preserving one dump entry per UID.
---
Nitpick comments:
In `@test/python/sdpa/mxfp8.py`:
- Around line 35-57: Review quantize_mxfp8_compact for a supported Transformer
Engine API for chunked MXFP8 quantization and concatenation, replacing direct
access to the result private attributes where possible. If no public API exists,
document or track this coupling explicitly while preserving the current
MXFP8Tensor reconstruction behavior.
🪄 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: 3e8026dc-e1b0-41f9-8626-fc403159de80
📒 Files selected for processing (5)
include/cudnn_frontend/graph_interface.hinclude/cudnn_frontend/node/scaled_dot_product_flash_attention.htest/python/sdpa/fp8.pytest/python/sdpa/mxfp8.pytest/python/test_mhas_v2.py
|
Evaluated the Transformer Engine coupling note. TE 2.12+ exposes |
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-386-848ad26 |
3 participants
@coderabbitai ignore
Why
These are the repro fixes shared by the UID (#280) and hybrid TID+UID (#383) graph-JSON designs. They do not depend on either identity architecture.
Two independent failures remained:
--perfrepros still built large references; large MXFP8 inputs also exceeded Transformer Engine's CUDA grid-y limit when quantized as one flattened tensorChanges
Runtime tensor dumps
SEQ_LEN_Q,SEQ_LEN_KV, and Q/K/V/O/dO/Stats/dQ/dK/dV ragged-offset tensorsFP8/MXFP8 repro testers
--perfScope
This PR does not change graph JSON, graph identity, tensor identity, extraction schema, or log format. #280 and #383 both depend on it.
When tensor-dump logging is disabled, the fast-path adapter returns before allocating its UID map. With logging enabled, existing dump behavior includes device-to-host copies and stream synchronization.
Verification
NVIDIA B200, cuDNN backend 9.30:
--perfpassedMXFP8 verification used the supported PyTorch 26.05 container with PyTorch 2.12.0a0, CUDA 13.2, cuDNN 9.30, and Transformer Engine 2.15. During the original investigation, these tester fixes moved the supplied set from 2/60 to 60/60 passes and the deduplicated raw-log set from 385/445 to 445/445 passes.
Summary by CodeRabbit
New Features
Performance
Tests