Skip to content

Commit 2f91701

Browse files
committed
[fix] remove incorrect V-row permutation in scaled_fp4_quant_trans_kernel
The Edenzzzz causal fix (PR #1455) applied a seq_k row permutation to V in scaled_fp4_quant_trans_kernel before the transpose step. This scrambles V rows after the transpose, causing the PV MMA to read values from wrong tokens and producing broken output for non-causal attention. The permutation is unnecessary: unlike K's N-dimension in the QK MMA, V's seq_k dimension is the K-dimension (reduction axis) of the PV MMA and is read sequentially. The transpose in scaled_fp4_quant_trans_kernel already produces the correct layout without any row reordering.
1 parent 7cebf5f commit 2f91701

1 file changed

Lines changed: 6 additions & 11 deletions

File tree

fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -286,24 +286,19 @@
286286

287287
// load input
288288
const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN;
289-
// Permute V rows within each 32-element block so the PV MMA K-indexed
290-
// access reads the correct CLayout N-indexed values (Edenzzzz causal fix).
291-
const int k_intra = token_id & 31;
292-
const int load_token_id = (token_id & ~31)
293-
| ((k_intra & 6) << 2) | ((k_intra & 24) >> 2) | (k_intra & 1);
294-
289+
295290
PackedVec in_vec;
296-
291+
297292
#pragma unroll
298293
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
299294
reinterpret_cast<uint32_t&>(in_vec.elts[i]) = 0;
300295
}
301-
302-
if (load_token_id < num_tokens) {
303-
in_vec = reinterpret_cast<PackedVec const*>(input +
296+
297+
if (token_id < num_tokens) {
298+
in_vec = reinterpret_cast<PackedVec const*>(input +
304299
batch_id * stride_bz_input + // batch dim
305300
head_id * stride_h_input + // head dim
306-
load_token_id * stride_seq_input + // seq dim (permuted)
301+
token_id * stride_seq_input + // seq dim
307302
(threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD)[0]; // feature dim
308303
}
309304

0 commit comments

Comments
 (0)