FP4 inputs of `dot_scaled` need more hints on documentation

[lightb0x]

Problem as-is

Documentation hides detail such that user need to find out failing cases by oneself.
Specifically, implementation vs. documentation as follows:

• dot_scaled requires rhs_k_pack=True for FP4 inputs
  • both for MXFP4 and NVFP4
  • for most of Blackwells (test required for GB100, GB200, etc.)
• documentation does not state that

Specifically

dot_scaled on CUDA with FP4 input require rhs_k_pack=True (the default value)
Current test cases clearly shows that:

• This case for NVFP4
• This case for MXFP4 and NVFP4 on SM120

Personally tested dot_scaled with MXFP4 input and rhs_k_pack=False.
It fails during MLIR passes with input shape error as follows:

python3: /project/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv2.cpp:760:
LogicalResult (anonymous namespace)::convertMMAImpl(
  DotOpInterface, Value, Value, Value,
  const LLVMType Converter *,
  ConversionPatternRewriter &,
  TensorCoreType,
  const NumRegisters &,
  const std::map<TensorCoreType, std::string> &,
  const EmitMmaCallback &
): Assertion `repA[2] == repB[1]' failed.

...

%148 = ttg.local_load %124 : !ttg.memdesc<128x64xi8, #shared, #smem> -> tensor<128x64xi8, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>>
%149 = ttg.local_load %144 : !ttg.memdesc<128x64xi8, #shared1, #smem> -> tensor<128x64xi8, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>
%150 = ttg.convert_layout %134 : tensor<128x4xi8, #blocked1> -> tensor<128x4xi8, #linear>
%151 = ttg.convert_layout %147 : tensor<128x4xi8, #blocked1> -> tensor<128x4xi8, #linear1>

%152 = tt.dot_scaled
  %148 scale %150,  # **lhs 128x64 i8 (128x128 fp4), scale 128x4**
  %149 scale %151,  # **rhs 128x64 i8 (128x128 fp4), scale 128x4**
  %arg14 lhs = e2m1 rhs = e2m1 {fastMath = false, **rhs_k_pack = false**}
:
tensor<128x64xi8, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>>,
tensor<128x4xi8, #linear>
*
tensor<128x64xi8, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>,
tensor<128x4xi8, #linear1>
-> tensor<128x128xf32, #mma>
scf.yield %152 : tensor<128x128xf32, #mma>

To-be

• documentation (live doc, code) need to clearly state conditions that rhs_k_pack=True is required.

---

Tested environment

• nvcr.io/nvidia/pytorch:25.10-py3 container with triton v3.6.0
• Blackwell GB203 (5070Ti)
• compute-capability=120 ptx-version=88
• nvidia-driver-open 580.126.09

I looked through all issues with dot_scaled, both open and closed.
Please correct me if I'm wrong.

[masahi]

I believe rhs_k_pack=False is actually supported on sm100. But, yes, on sm120 it is not. Feel free to update the documentation.