VK_NV_cooperative_matrix2: coopMatLoadTensorNV with tensorViewTranspose and decode callback provides incorrect blockCoords/coordInBlock

Gaming and Visualization Technologies APIs Vulkan
,
1 
Hi NVIDIA Vulkan team,

I have a minimal reproducer for what appears to be an incorrect coordinate mapping issue in `VK_NV_cooperative_matrix2` when using:

```glsl
coopMatLoadTensorNV(mat, data, elementOffset, layout, tensorViewTranspose, decodeFunc)

Summary

The tested input is a Q4_K quantized block matrix whose physical layout stores K blocks first:

[K_block, M]

For a logical matrix element (m, k), the expected Q4_K block location is:

k_block = k / QUANT_K
idx     = k % QUANT_K
block   = k_block * M + m

The reproducer directly constructs this [K_block, M] input buffer. Therefore the test focuses only on the CM2 read path: whether coopMatLoadTensorNV + tensorViewTranspose + decode callback supplies correct blockCoords and coordInBlock values to the decode callback.

Expected behavior

For each cooperative-matrix component, the hardware path should match the software reference path:

HW path: coopMatLoadTensorNV + tensorViewTranspose + decode callback
SW path: coopMatPerElementNV(row, col) + manual Q4_K address calculation

Expected: hw_val == sw_val

The decode callback expects coordinates equivalent to:

blockCoords[0]  = k_block
blockCoords[1]  = m row
coordInBlock[0] = offset within the Q4_K block along K, range 0..255
coordInBlock[1] = 0

Actual behavior

On my system, the HW path produces many mismatches against the SW reference path.

Environment:

GPU:        NVIDIA GeForce RTX 4080
Driver:     610.47.0.0
Vulkan SDK: 1.4.341.1
OS:         Windows

Example result:

RESULT: *** FAIL ***
Total mismatches: 63056
block_k=0      : 15928 mismatches <-- FAIL
block_k=128    : 15376 mismatches <-- FAIL
block_k=256    : 15368 mismatches <-- FAIL
block_k=384    : 16384 mismatches <-- FAIL

Example mismatch detail:

tid=128 block_k=0 comp=0
  HW: blockCoords=(0,0) coordInBlock=(0,16) read block=0 idx=0 val=0.0
  SW: row=16 col=0 k_block=0 m=16 read block=16 idx=0 val=3.0 delta=-3.0

For the same cooperative-matrix component, the SW reference expects row=16, col=0, therefore it reads block=16, idx=0. However, the HW decode callback receives blockCoords=(0,0), coordInBlock=(0,16), which makes it read block=0, idx=0.

This suggests that the callback coordinate mapping does not match the expected tensor-view interpretation. Some runs also show invalid-looking coordinate values such as 0xffffffff, and earlier experiments found that accessing the callback buffer-reference parameter can cause VK_ERROR_DEVICE_LOST.

Reproducer

I attached a zip file containing:

nv_cm2_transpose_bug.comp
nv_cm2_transpose_bug.cpp
nv_cm2_transpose_bug.spv
nv_cm2_transpose_bug_report_en.md

Build commands:

glslc --target-env=vulkan1.3 -o nv_cm2_transpose_bug.spv nv_cm2_transpose_bug.comp

cl /EHsc /std:c++17 /I"C:\VulkanSDK\1.4.341.1\Include" nv_cm2_transpose_bug.cpp ^
  /link /LIBPATH:"C:\VulkanSDK\1.4.341.1\Lib" vulkan-1.lib

Run:

nv_cm2_transpose_bug.exe

Request

Could NVIDIA please confirm whether this is a driver bug in VK_NV_cooperative_matrix2 / GL_NV_cooperative_matrix2, specifically for this combination?

col-major [K_block, M] tensor layout
+ tensorViewNV transpose
+ block-quantized decode callback

Expected fix/behavior:

  1. The decode callback should receive spec-compliant blockCoords and coordInBlock values when the tensor layout uses block size, such as Q4_K QUANT_K=256.
  2. Coordinate mapping should remain correct when tensorViewNV<2, false, 1, 0> transpose view is used.
  3. The callback buffer-reference parameter should point to a valid block, and accessing its fields should not cause VK_ERROR_DEVICE_LOST.
  4. After the fix, this reproducer should print RESULT: PASS.

Thanks.

nv_cm2_transpose_bug_issue_package.zip (18.7 KB)

2

Hi,

I think this line:

tl = setTensorLayoutBlockSizeNV(tl, 0, QUANT_K);   // dim 0 (K) blocked

should be:

tl = setTensorLayoutBlockSizeNV(tl, QUANT_K, 1);

i.e. the blocking is in the K dimension. With this change, the tests passes for me. Let me know what you think.