[milakov]

Description

Row-scaled NVFP4 activations currently go through per-expert quantization in the split path. Quantize eligible concatenated activations once, then partition their payloads and row-scale metadata into per-expert storage to reduce repeated kernel launches and setup overhead.

In the tested B300 split-quantization microbenchmarks, median synchronized operation latency decreased by 1.3–62.3%, depending on expert count, width, and split alignment. MXFP8 controls stayed within 1% of baseline. These are operation-level measurements, not full-model iteration timings.

Type of change

• Code refactoring

Changes

• Add a guarded fast path for contiguous BF16, rowwise-only, deterministic row-scaled NVFP4 activation quantization using existing native kernels.
• Preserve independent per-expert scale padding, including uneven and empty splits. Requantizing one expert must not modify another expert's storage.
• Retain the existing path for unsupported layouts, quantizer options, and allocation requirements.
• Add exact regression coverage for quantized payloads, scales, dequantized values, graph replay, expert requantization, hybrid quantizers, GroupedLinear outputs/gradients, CPU offload, debug dispatch, and GroupedTensor fallback.

This change is independent of #3799.

Validation

Validation results below were recorded against base 63451d6:

• 154 focused correctness and integration checks passed; 24 cases were skipped because UE5M3 requires compute capability 10.7.
• Full PyTorch debug suite: 18,806 passed.
• Full PyTorch unit suite: 62,118 passed, with the same 38 failures reproduced on unmodified main and no additional failures or errors.
• Formatting, lint, and license checks passed.

The baseline failures concern MXFP8 signed-zero comparisons and FlexAttention/cuDNN Frontend tests.

Merge compatibility was checked against upstream main at b021980; none of the five changed files has intervening upstream changes.

Checklist

• I have read the contributing guidelines.
• The functionality is complete.
• I have commented the code where needed.
• I have documented the new private helpers and eligibility constraints in the code.
• My changes generate no new warnings.
• I have added tests covering the change.
• New and existing unit tests pass locally with my changes (baseline failures documented above).

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable defects were found.

Summary

Adds a guarded fast path that quantizes eligible row-scaled NVFP4 activations once, then splits their packed values and scale data by expert.

• GroupedLinear quantizes eligible row-scaled NVFP4 activations once.

No actionable issues were found. Tests were inspected but not run during this review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Split quantization request] --> B{Eligible row-scaled NVFP4?}
  B -->|No| C[Existing dispatch]
  B -->|Yes, bulk allocation allowed| D[Quantize concatenated input once]
  D --> E[Split packed values and row amax]
  D --> F[Split scales with separate expert padding]
  E --> G[Per-expert storage]
  F --> G

Loading

_{Reviews (1) · Last reviewed commit: "Reduce NVFP4 activation quantization ove..." · Reviewed by Greptile}