Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughMLA now supports optional per-token or per-group FP8 CKV scales and zero-dimensional KPE layouts. The wrapper validates and forwards scale arrays. CUDA bindings and kernels apply the scales during FP8 repacking. Trace schemas and tests cover FA2 and FA3 paths. ChangesFP8 CKV scaling and zero-KPE MLA
Estimated code review effort: 4 (Complex) | ~60 minutes Mergeability Score: ⚪ Minimal · up to The change adds FP8 KV and zero-width positional-key support with validation and targeted GPU coverage; no actionable merge-blocking risk remains beyond normal checks and review. Sequence Diagram(s)sequenceDiagram
participant MLAWrapper
participant CUDABinding
participant MLAParams
participant FP8MLAKernel
MLAWrapper->>CUDABinding: pass validated ckv_scale_arr
CUDABinding->>MLAParams: store scale-array pointer
MLAParams->>FP8MLAKernel: provide CKV scale metadata
FP8MLAKernel->>FP8MLAKernel: repack FP8 CKV to BF16
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
[FAILED] Pipeline #62257276 — 10/18 executed test jobs passed Compared with nightly #62109159. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsNew relative to nightly (attribution uncertain)
Pre-existing failures
|
JustinTong0323
changed the title
|
I added the requested SM90 FP8 KV NoPE support, including optional per-token/per-128-channel CKV scales for both FA2 and FA3, and rebased the branch onto the latest main. On an H200 with CUDA 12.8 / PyTorch 2.11.0+cu128, the new non-identity-page-table group-scale test passes on both FA2 and FA3; the existing KPE=64 per-tensor FP8 and row-aliasing regressions also pass. The current required check is red only because CI was skipped pending authorization for the new commit ( |
![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
🧹 Nitpick comments (3)
flashinfer/mla/_core.py (1)
2523-2523: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse tuple unpacking instead of tuple concatenation.
Ruff reports RUF005 on this line.
♻️ Proposed change
- expected_scale_shape = ckv_cache.shape[:-1] + (self._head_dim_ckv // 128,) + expected_scale_shape = (*ckv_cache.shape[:-1], self._head_dim_ckv // 128)🤖 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 `@flashinfer/mla/_core.py` at line 2523, Update the expected_scale_shape assignment to construct the shape tuple using tuple unpacking instead of concatenating ckv_cache.shape[:-1] with a one-element tuple, while preserving the same resulting dimensions.Source: Linters/SAST tools
include/flashinfer/attention/mla_hopper.cuh (1)
392-409: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMerge the duplicated scale-array branch.
Lines 393 and 399 test
ckv_scale_arrtwice in a row. Fold the lookup and the scaled conversion into one branch. Add a short comment that states why the scale-array path uses scalar FP32 multiplication instead of__hmul2.🤖 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 `@include/flashinfer/attention/mla_hopper.cuh` around lines 392 - 409, The scale-array handling in the surrounding attention conversion code duplicates the ckv_scale_arr branch. Merge the scale lookup and conversion loop into a single ckv_scale_arr branch, retaining scalar FP32 multiplication there and the existing packed __hmul2 path in the else branch. Add a brief comment explaining why the scale-array path requires scalar FP32 multiplication.include/flashinfer/attention/mla.cuh (1)
296-348: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated CKV scale-lookup branch in
repack_fp8_kv_to_bf16. Both kernels copy the same block: a firstif (ckv_scale_arr != nullptr)that resolves the physical row and reads the group scale, then a second identical test that selects scalar FP32 multiplication or packed__hmul2. The shared root cause is the copied scale-lookup logic. Extract one helper that returns the CKV scale for a(row, col)pair and applies it, then call it from both kernels.
include/flashinfer/attention/mla.cuh#L296-L348: merge the twockv_scale_arr != nullptrtests at Lines 323 and 331 into one branch, and add the hot-path rationale comment for the FP8-to-BF16 staging and the scalar-scale path.include/flashinfer/attention/mla_hopper.cuh#L392-L409: merge the twockv_scale_arrtests at Lines 393 and 399 into the same shared helper.🤖 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 `@include/flashinfer/attention/mla.cuh` around lines 296 - 348, Extract the duplicated CKV scale lookup and application logic into one shared helper for `(row, col)` pairs, preserving paged `physical_row` resolution and both scalar and packed multiplication paths; invoke it from `repack_fp8_kv_to_bf16` in include/flashinfer/attention/mla.cuh:296-348, merging the two `ckv_scale_arr` branches and adding the requested hot-path rationale comment. Apply the same helper and branch consolidation in include/flashinfer/attention/mla_hopper.cuh:392-409.
🤖 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 `@flashinfer/mla/_core.py`:
- Around line 2518-2532: Add ckv_scale_arr as an optional tensor in
mla_paged_decode_trace, using the existing shape, dtype, device, and contiguity
validation behavior. Update the trace template and regenerate both generated
fixtures so they include this optional tensor and preserve the Invalid shape of
ckv_scale_arr validation prefix.
In `@include/flashinfer/attention/mla.cuh`:
- Around line 350-373: The FP8 KPE repacking path in
load_kv/repack_fp8_kv_to_bf16 uses the colliding standard KPE swizzle for FA2
head_dim_kpe=64. Mirror SWIZZLE_MODE_KPE_RAW, LANES_PER_ROW_KPE, and
INNER_LOADS_KPE from mla_hopper.cuh, use the raw swizzle when reading KPE, and
add FA2 coverage for head_dim_kpe=64.
---
Nitpick comments:
In `@flashinfer/mla/_core.py`:
- Line 2523: Update the expected_scale_shape assignment to construct the shape
tuple using tuple unpacking instead of concatenating ckv_cache.shape[:-1] with a
one-element tuple, while preserving the same resulting dimensions.
In `@include/flashinfer/attention/mla_hopper.cuh`:
- Around line 392-409: The scale-array handling in the surrounding attention
conversion code duplicates the ckv_scale_arr branch. Merge the scale lookup and
conversion loop into a single ckv_scale_arr branch, retaining scalar FP32
multiplication there and the existing packed __hmul2 path in the else branch.
Add a brief comment explaining why the scale-array path requires scalar FP32
multiplication.
In `@include/flashinfer/attention/mla.cuh`:
- Around line 296-348: Extract the duplicated CKV scale lookup and application
logic into one shared helper for `(row, col)` pairs, preserving paged
`physical_row` resolution and both scalar and packed multiplication paths;
invoke it from `repack_fp8_kv_to_bf16` in
include/flashinfer/attention/mla.cuh:296-348, merging the two `ckv_scale_arr`
branches and adding the requested hot-path rationale comment. Apply the same
helper and branch consolidation in
include/flashinfer/attention/mla_hopper.cuh:392-409.
🪄 Autofix
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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 20e6738e-cbe3-4e92-982d-c77a728082f5
📒 Files selected for processing (9)
csrc/batch_mla_binding.cucsrc/batch_mla_run.cucsrc/batch_mla_sm90_binding.cucsrc/batch_mla_sm90_run.cuflashinfer/mla/_core.pyinclude/flashinfer/attention/mla.cuhinclude/flashinfer/attention/mla_hopper.cuhinclude/flashinfer/attention/mla_params.cuhtests/attention/test_deepseek_mla.py
|
/bot run tests/attention |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/trace/templates/attention.py (1)
1428-1437: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winAdd FP8 scale support to
mla_paged_decodeinit and reference functions.The trace schema declares optional FP8 inputs (
ckv_scale_arr,ckv_scale,kpe_scale), but_mla_paged_decode_initdiscardsckv_scale_groupswithout constructing scale tensors, and_mla_paged_decode_referencedoes not accept these scale parameters. If trace replay forwards these optional inputs, the reference invocation fails. Either implement FP8 dequantization in the reference and scale tensor construction in the init, or move FP8 variants to a separate trace template and mark this template for non-quantized KV only.🤖 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 `@flashinfer/trace/templates/attention.py` around lines 1428 - 1437, The MLA paged decode trace template must handle the schema’s optional FP8 inputs consistently. Update _mla_paged_decode_init to construct the required scale tensors from ckv_scale_groups, and extend _mla_paged_decode_reference and its invocation to accept and apply ckv_scale_arr, ckv_scale, and kpe_scale; alternatively, restrict this template to non-quantized KV and route FP8 variants through a separate template.
🧹 Nitpick comments (1)
tests/trace/test_fi_trace.py (1)
683-703: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAssert the complete
ckv_scale_arrcontract.The test checks only the symbolic shape. It does not check that the emitted input remains optional and uses
float32. Add assertions fordefn["inputs"]["ckv_scale_arr"]["optional"] is Trueanddefn["inputs"]["ckv_scale_arr"]["dtype"] == "float32".🤖 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/trace/test_fi_trace.py` around lines 683 - 703, Update the assertions for ckv_scale_arr in the BatchMLAPagedAttentionWrapper trace test to also require optional == True and dtype == "float32", while preserving the existing symbolic shape assertion.
🤖 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.
Outside diff comments:
In `@flashinfer/trace/templates/attention.py`:
- Around line 1428-1437: The MLA paged decode trace template must handle the
schema’s optional FP8 inputs consistently. Update _mla_paged_decode_init to
construct the required scale tensors from ckv_scale_groups, and extend
_mla_paged_decode_reference and its invocation to accept and apply
ckv_scale_arr, ckv_scale, and kpe_scale; alternatively, restrict this template
to non-quantized KV and route FP8 variants through a separate template.
---
Nitpick comments:
In `@tests/trace/test_fi_trace.py`:
- Around line 683-703: Update the assertions for ckv_scale_arr in the
BatchMLAPagedAttentionWrapper trace test to also require optional == True and
dtype == "float32", while preserving the existing symbolic shape assertion.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 3ca1ce68-c82a-48d7-b479-312419b21510
📒 Files selected for processing (8)
flashinfer/mla/_core.pyflashinfer/trace/templates/attention.pyinclude/flashinfer/attention/mla.cuhinclude/flashinfer/attention/mla_hopper.cuhtests/attention/test_deepseek_mla.pytests/trace/fi_trace_out/mla_paged_decode_h16_ckv512_kpe64_ps1.jsontests/trace/fi_trace_out/mla_paged_decode_h16_ckv512_kpe64_ps64.jsontests/trace/test_fi_trace.py
🚧 Files skipped from review as they are similar to previous changes (4)
- flashinfer/mla/_core.py
- tests/attention/test_deepseek_mla.py
- include/flashinfer/attention/mla_hopper.cuh
- include/flashinfer/attention/mla.cuh
|
/bot run tests/attention |
|
[FAILED] Pipeline #62369148 — 12/18 executed test jobs passed Compared with nightly #62109159. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsPre-existing failures
|
|
/bot run tests/attention |
saltyminty
left a comment
saltyminty
left a comment
There was a problem hiding this comment.
Left one non-blocking comment.
| *, | ||
| ckv_scale: Optional[float] = None, | ||
| kpe_scale: Optional[float] = None, | ||
| ckv_scale_arr: Optional[torch.Tensor] = None, |
There was a problem hiding this comment.
Non-blocking but just want to hear thoughts: the intention is that we will accept only one of either ckv_scale or ckv_scale_arr, correct?
- Since these arguments are already keyword only, we are free to reorder them (i.e. putting
ckv_scale_arrdirectly afterckv_scale) - would it be cleaner to only keep one argument as
ckv_scale: Union[float, torch.Tensor]?
There was a problem hiding this comment.
Yes, they are intended to be mutually exclusive. I reordered ckv_scale_arr directly after ckv_scale and now require exactly one of them for FP8 KV.
I kept them as separate arguments because the trace schema represents scalar and tensor inputs with different descriptors. A single Union[float, torch.Tensor] parameter cannot faithfully describe both forms in the generated benchmark definition. Internally, the grouped-scale path still passes 1.0 through the unused scalar launcher slot, so the C++ ABI remains unchanged.
|
[SUCCESS] Pipeline #62545365: 18/18 executed test jobs passed |
🚨 POTENTIAL BREAKING PUBLIC API CHANGE DETECTED 🚨Caution THIS PR APPEARS TO BREAK THE PUBLIC API. AUTHORS AND REVIEWERS: DO NOT MISS THIS. This is an advisory warning and does not gate merging. Confirm compatibility and provide a deprecation or migration path, or track the fix in a follow-up PR. 1 public API finding(s):
|
|
/bot run tests/attention |
|
Regarding the public API advisory above: this change is backward-compatible. |
|
@flashinfer-bot run |
saltyminty
left a comment
saltyminty
left a comment
There was a problem hiding this comment.
Supposed API breaking change is in a keyword only section and is not an issue.
## Description
Allow `BatchMLAPagedAttentionWrapper` to run with `head_dim_kpe=0`,
including FP8 KV caches on SM90.
When the positional-key-embedding dimension is zero, the KPE MMA loop is
empty. The FA2 and FA3 paths initialize the QK accumulator from
compressed KV instead. For FP8 KV, both native backends repack to BF16
and can consume an optional contiguous FP32 scale tensor with one scale
per physical KV token and 128 CKV channels.
This change:
- supports zero-width KPE in FA2 and FA3;
- supports FP8 E4M3 KV with BF16 queries on SM90 for `head_dim_ckv=512`
and `head_dim_kpe` in `{0, 64}`;
- adds optional `ckv_scale_arr` with shape `ckv_cache.shape[:-1] + (4,)`
for per-token, per-group CKV dequantization;
- preserves existing per-tensor `ckv_scale` / `kpe_scale` behavior;
- validates backend, device, dtype, dimensions, scale shape, dtype,
device, and contiguity;
- adds paged-MLA correctness coverage with a non-identity page table for
FA2 and FA3.
## Related Issues
None.
## Pull Request Checklist
- [x] Pre-commit hooks pass on the changed files.
- [x] Tests have been added or updated as needed.
- [x] GPU tests pass.
## Tests
NVIDIA H200, CUDA 12.8, PyTorch 2.11.0+cu128:
- `test_batch_mla_without_kpe`: FA2 and FA3 passed.
- `test_batch_mla_fp8_nope_group_scales_matches_bf16_reference`: FA2 and
FA3 passed.
- `test_fp8_kv_scales_are_keyword_only` and
`test_fp8_kv_requires_scales`: passed.
- Representative existing KPE=64 per-tensor FP8 and row-aliasing
regressions: passed.
## Reviewer Notes
- `ckv_scale_arr` follows the physical paged-cache layout used by the
downstream SGLang integration: `[num_pages, page_size, head_dim_ckv /
128]`.
- Existing BF16/FP16 and FP8 per-tensor scale behavior is unchanged.
- The latest external attention CI failures before this update were not
stable across reruns: HCA failures were classified as pre-existing,
while the new failures alternated between GB200 multi-node NCCL setup
and RTX Pro 6000 `cudaErrorUnknown` failures.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added FP8 batched paged attention with optional per-token/per-channel
CKV scaling.
* Added support for FA2 and FA3 FP8 execution.
* Expanded supported configurations to include zero- or 64-dimensional
key positional encoding.
* **Bug Fixes**
* Fixed attention computation without key positional encoding
dimensions.
* Added validation for incompatible scale inputs, shapes, layouts, and
configurations.
* **Tests**
* Added coverage for grouped CKV scaling and trace inputs.
* Added BF16 reference comparisons across supported implementations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Co-authored-by: Lee Nau <lnau@nvidia.com>
4 participants
Description
Allow
BatchMLAPagedAttentionWrapperto run withhead_dim_kpe=0, including FP8 KV caches on SM90.When the positional-key-embedding dimension is zero, the KPE MMA loop is empty. The FA2 and FA3 paths initialize the QK accumulator from compressed KV instead. For FP8 KV, both native backends repack to BF16 and can consume an optional contiguous FP32 scale tensor with one scale per physical KV token and 128 CKV channels.
This change:
head_dim_ckv=512andhead_dim_kpein{0, 64};ckv_scale_arrwith shapeckv_cache.shape[:-1] + (4,)for per-token, per-group CKV dequantization;ckv_scale/kpe_scalebehavior;Related Issues
None.
Pull Request Checklist
Tests
NVIDIA H200, CUDA 12.8, PyTorch 2.11.0+cu128:
test_batch_mla_without_kpe: FA2 and FA3 passed.test_batch_mla_fp8_nope_group_scales_matches_bf16_reference: FA2 and FA3 passed.test_fp8_kv_scales_are_keyword_onlyandtest_fp8_kv_requires_scales: passed.Reviewer Notes
ckv_scale_arrfollows the physical paged-cache layout used by the downstream SGLang integration:[num_pages, page_size, head_dim_ckv / 128].cudaErrorUnknownfailures.Summary by CodeRabbit
New Features
Bug Fixes
Tests