[Bugfix][ROCm] Mask paged attention V cache padding (#53856)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Sep 8, 2026 at 14:39 UTC
755541443f62a1f305199b1c82e524abccec7e13
2 files changed
+222
-41
csrc/rocm/attention.cu
+196
-31
@@ -59,6 +59,37 @@ enum class MFMAType {
59
Fp4 = 2,
60
};
61
62
+// Score masking is insufficient for unwritten V padding: 0 * NaN is NaN.
63
+template <typename cache_t, typename Vec>
64
+__device__ __forceinline__ Vec mask_v_cache_padding(Vec values,
65
+ const int first_token_idx,
66
+ const int seq_len) {
67
+ constexpr int NUM_ELEMS = sizeof(Vec) / sizeof(cache_t);
68
+ if (first_token_idx + NUM_ELEMS <= seq_len) {
69
+ return values;
70
+ }
71
+ if (first_token_idx >= seq_len) {
72
+ return Vec{};
73
+ }
74
+ constexpr int WORD_BYTES = sizeof(uint64_t);
75
+ constexpr int NUM_WORDS = sizeof(Vec) / WORD_BYTES;
76
+ static_assert(sizeof(Vec) % WORD_BYTES == 0);
77
+ const int valid_bytes = (seq_len - first_token_idx) * sizeof(cache_t);
78
+ uint64_t* words = reinterpret_cast<uint64_t*>(&values);
79
+#pragma unroll
80
+ for (int i = 0; i < NUM_WORDS; i++) {
81
+ const int word_valid_bytes = valid_bytes - i * WORD_BYTES;
82
+ const uint64_t word_mask =
83
+ word_valid_bytes >= WORD_BYTES
84
+ ? UINT64_MAX
85
+ : (word_valid_bytes <= 0
86
+ ? 0
87
+ : (uint64_t{1} << (word_valid_bytes * 8)) - 1);
88
+ words[i] &= word_mask;
89
+ }
90
+ return values;
91
+}
92
+
93
#if defined(__HIP__GFX9__)
94
95
#define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32
@@ -82,6 +113,30 @@ typedef struct _B16x8 {
113
_B16x4 xy[2];
114
} _B16x8;
115
116
+__device__ __forceinline__ _B16x8
117
+mask_16b_v_cache_boundary(_B16x8 values, const int valid_bytes) {
118
+ // The caller guarantees 0 < valid_bytes < 16.
119
+ uint64_t* words = reinterpret_cast<uint64_t*>(&values);
120
+ if (valid_bytes <= 8) {
121
+ const uint64_t low_mask =
122
+ valid_bytes == 8 ? UINT64_MAX : (uint64_t{1} << (valid_bytes * 8)) - 1;
123
+ words[0] &= low_mask;
124
+ words[1] = 0;
125
+ } else {
126
+ const uint64_t high_mask = (uint64_t{1} << ((valid_bytes - 8) * 8)) - 1;
127
+ words[1] &= high_mask;
128
+ }
129
+ return values;
130
+}
131
+
132
+__device__ __forceinline__ uint2
133
+mask_8b_v_cache_boundary(uint2 values, const int valid_bytes) {
134
+ // The caller guarantees 0 < valid_bytes < 8.
135
+ uint64_t* word = reinterpret_cast<uint64_t*>(&values);
136
+ *word &= (uint64_t{1} << (valid_bytes * 8)) - 1;
137
+ return values;
138
+}
139
+
140
using _B8x8 = uint2;
141
using _B8x4 = int32_t; // used in builtins
142
using bit8_t = uint8_t;
@@ -318,7 +373,7 @@ template <typename scalar_t, typename cache_t,
373
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
374
int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE>
375
__global__
321
-__launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
376
+__launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel(
377
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
378
const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
379
const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
@@ -789,29 +844,72 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
844
_B16x4 outelems[VHELOOP];
845
// Softmax V mfma
846
// v layout: 16he across lanes x 16 tokens per lane
792
- for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
793
- floatx4 tmp_out = {0};
794
-
795
- for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
796
- if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
847
+ if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
848
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
849
+ floatx4 tmp_out = {0};
850
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
851
for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
852
+ if (__builtin_expect(partition_start_token_idx + T_PAR_SIZE > seq_len,
853
+ 0)) {
854
+ const int vglobal_fetch_start =
855
+ partition_start_token_idx +
856
+ vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP +
857
+ rowid * VTOKENS_PER_LANE +
858
+ vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
859
+ if (vglobal_fetch_start >= seq_len) {
860
+ continue;
861
+ }
862
+ if (vglobal_fetch_start + CONTIGUOUS_KV_ELEMS_16B_LOAD > seq_len) {
863
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth] =
864
+ mask_16b_v_cache_boundary(
865
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth],
866
+ (seq_len - vglobal_fetch_start) * sizeof(cache_t));
867
+ }
868
+ }
869
for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) {
870
const int offset = rowid * VTLANELOOP * ELEMS8_ELEMS4_RATIO +
871
vfetch_depth * ELEMS8_ELEMS4_RATIO + i;
872
const int offset1 = offset % ROWS_PER_WARP;
873
const int offset2 = offset / ROWS_PER_WARP;
803
- // output format is 16 qheads across 16 lanes, 16 head elems spread
804
- // across 4 rows
874
tmp_out = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
875
Vlocal[vtoken_depth][vhe_depth][vfetch_depth].xy[i],
876
shared_logits[vtoken_depth][offset2][lane16id][offset1],
877
tmp_out);
878
}
879
}
811
- // KV cache fp8
812
- } else {
880
+ }
881
+ outelems[vhe_depth] = from_floatx4<scalar_t>(tmp_out);
882
+ }
883
+ } else {
884
+ const int valid_vtoken_depths =
885
+ min(VTLOOP, DIVIDE_ROUND_UP(seq_len - partition_start_token_idx,
886
+ VTOKENS_PER_LANE * ROWS_PER_WARP));
887
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
888
+ floatx4 tmp_out = {0};
889
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
890
for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
891
_B16x8 Vtmp = Vlocal[vtoken_depth][vhe_depth][vfetch_depth];
892
+ if (__builtin_expect(partition_start_token_idx + T_PAR_SIZE > seq_len,
893
+ 0)) {
894
+ if (vtoken_depth >= valid_vtoken_depths) {
895
+ continue;
896
+ }
897
+ if (vtoken_depth + 1 == valid_vtoken_depths) {
898
+ const int vglobal_fetch_start =
899
+ partition_start_token_idx +
900
+ vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP +
901
+ rowid * VTOKENS_PER_LANE +
902
+ vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
903
+ if (vglobal_fetch_start >= seq_len) {
904
+ continue;
905
+ }
906
+ if (vglobal_fetch_start + CONTIGUOUS_KV_ELEMS_16B_LOAD >
907
+ seq_len) {
908
+ Vtmp = mask_16b_v_cache_boundary(Vtmp,
909
+ seq_len - vglobal_fetch_start);
910
+ }
911
+ }
912
+ }
913
// reinterpret V format as 16 elements of 8bits
914
_B8x16 Vtmp8x16 = *reinterpret_cast<_B8x16*>(&Vtmp);
915
for (int j = 0; j < ELEMS16_ELEMS8_RATIO; j++) {
@@ -856,12 +954,9 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
954
}
955
}
956
}
859
- }
860
- // apply post Softmax V mfma v_scale
861
- if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
957
tmp_out *= *v_scale;
958
+ outelems[vhe_depth] = from_floatx4<scalar_t>(tmp_out);
959
}
864
- outelems[vhe_depth] = from_floatx4<scalar_t>(tmp_out);
960
}
961
962
__syncthreads();
@@ -1334,15 +1429,37 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1429
}
1430
}
1431
} else { // warp in context
1432
+ #define SV_mfma_value(x, value) \
1433
+ for (int qh = 0; qh < QHLOOP; qh++) { \
1434
+ acc[qh] = gcn_mfma4x4x4_instr<scalar_t, 4, 2 * x, 0>( \
1435
+ logits[qh], value.xy[0], acc[qh]); \
1436
+ acc[qh] = gcn_mfma4x4x4_instr<scalar_t, 4, 2 * x + 1, 0>( \
1437
+ logits[qh], value.xy[1], acc[qh]); \
1438
+ }
1439
+
1440
#define SV_mfma(x) \
1441
if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \
1442
Vlocal[vh][x] = convert_b8x8_custom<scalar_t>(Vlocalb8[vh][x]); \
1443
} \
1341
- for (int qh = 0; qh < QHLOOP; qh++) { \
1342
- acc[qh] = gcn_mfma4x4x4_instr<scalar_t, 4, 2 * x, 0>( \
1343
- logits[qh], Vlocal[vh][x].xy[0], acc[qh]); \
1344
- acc[qh] = gcn_mfma4x4x4_instr<scalar_t, 4, 2 * x + 1, 0>( \
1345
- logits[qh], Vlocal[vh][x].xy[1], acc[qh]); \
1444
+ SV_mfma_value(x, Vlocal[vh][x])
1445
+
1446
+ #define SV_mfma_tail(x) \
1447
+ if (x * 8 < valid_warp_tokens) { \
1448
+ _B16x8 Vtmp; \
1449
+ if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \
1450
+ Vtmp = Vlocal[vh][x]; \
1451
+ if ((x + 1) * 8 > valid_warp_tokens) { \
1452
+ Vtmp = mask_16b_v_cache_boundary( \
1453
+ Vtmp, (valid_warp_tokens - x * 8) * sizeof(cache_t)); \
1454
+ } \
1455
+ } else { \
1456
+ _B8x8 Vtmp8 = Vlocalb8[vh][x]; \
1457
+ if ((x + 1) * 8 > valid_warp_tokens) { \
1458
+ Vtmp8 = mask_8b_v_cache_boundary(Vtmp8, valid_warp_tokens - x * 8); \
1459
+ } \
1460
+ Vtmp = convert_b8x8_custom<scalar_t>(Vtmp8); \
1461
+ } \
1462
+ SV_mfma_value(x, Vtmp) \
1463
}
1464
1465
for (int vh = 0; vh < VHELOOP; vh++) {
@@ -1350,18 +1467,28 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1467
for (int qh = 0; qh < QHLOOP; qh++) {
1468
acc[qh] = {0};
1469
}
1353
- // SoftMax-V calculation
1354
- // logits -> token dimension is distributed across lanes
1355
- // Vlocal -> token dimension is depthwise within lane
1356
- // uses mfma instruction block broadcast for logits
1357
- SV_mfma(0);
1358
- SV_mfma(1);
1359
- SV_mfma(2);
1360
- SV_mfma(3);
1361
- SV_mfma(4);
1362
- SV_mfma(5);
1363
- SV_mfma(6);
1364
- SV_mfma(7);
1470
+ // SoftMax-V calculation. Skip fully invalid V vectors in the partial
1471
+ // tail warp, and only mask the one vector crossing seq_len.
1472
+ if (__builtin_expect(warp_start_token_idx + WARP_SIZE > seq_len, 0)) {
1473
+ const int valid_warp_tokens = seq_len - warp_start_token_idx;
1474
+ SV_mfma_tail(0);
1475
+ SV_mfma_tail(1);
1476
+ SV_mfma_tail(2);
1477
+ SV_mfma_tail(3);
1478
+ SV_mfma_tail(4);
1479
+ SV_mfma_tail(5);
1480
+ SV_mfma_tail(6);
1481
+ SV_mfma_tail(7);
1482
+ } else {
1483
+ SV_mfma(0);
1484
+ SV_mfma(1);
1485
+ SV_mfma(2);
1486
+ SV_mfma(3);
1487
+ SV_mfma(4);
1488
+ SV_mfma(5);
1489
+ SV_mfma(6);
1490
+ SV_mfma(7);
1491
+ }
1492
1493
for (int qh = 0; qh < QHLOOP; qh++) {
1494
if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
@@ -1373,6 +1500,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1500
}
1501
1502
#undef SV_mfma
1503
+ #undef SV_mfma_tail
1504
+ #undef SV_mfma_value
1505
} // warp in context
1506
1507
__syncthreads();
@@ -1963,6 +2092,24 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2092
}
2093
}
2094
2095
+ if (partition_start_token_idx + T_PAR_SIZE > seq_len) {
2096
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
2097
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
2098
+ for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
2099
+ const int vglobal_fetch_start =
2100
+ partition_start_token_idx +
2101
+ vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP +
2102
+ vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
2103
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth / 2]
2104
+ .xy[vfetch_depth % 2] = mask_v_cache_padding<cache_t>(
2105
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth / 2]
2106
+ .xy[vfetch_depth % 2],
2107
+ vglobal_fetch_start, seq_len);
2108
+ }
2109
+ }
2110
+ }
2111
+ }
2112
+
2113
floatx8 dout[TLOOP];
2114
// qk wmma
2115
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
@@ -2741,6 +2888,24 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2888
}
2889
}
2890
2891
+ if (partition_start_token_idx + T_PAR_SIZE > seq_len) {
2892
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
2893
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
2894
+ for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
2895
+ const int vglobal_fetch_start =
2896
+ partition_start_token_idx +
2897
+ vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP +
2898
+ rowid * VTOKENS_PER_LANE +
2899
+ vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
2900
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth] =
2901
+ mask_v_cache_padding<cache_t>(
2902
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth],
2903
+ vglobal_fetch_start, seq_len);
2904
+ }
2905
+ }
2906
+ }
2907
+ }
2908
+
2909
floatx8 dout[TLOOP];
2910
// qk wmma
2911
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
tests/kernels/attention/test_attention.py
+26
-10
@@ -153,18 +153,22 @@ def test_paged_attention(
153
if use_alibi:
154
alibi_slopes = torch.randn(num_query_heads, dtype=torch.float)
155
156
- seq_lens = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
157
- seq_lens[-1] = MAX_SEQ_LEN
158
- max_seq_len = max(seq_lens)
159
- seq_lens = torch.tensor(seq_lens, dtype=torch.int)
156
+ seq_lens_lst = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
157
+ seq_lens_lst[-1] = MAX_SEQ_LEN
158
+ max_seq_len = max(seq_lens_lst)
159
+ seq_lens = torch.tensor(seq_lens_lst, dtype=torch.int)
160
161
# Create the block tables.
162
max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
163
block_tables_lst: list[list[int]] = []
164
- for _ in range(num_seqs):
164
+ for seq_idx, seq_len in enumerate(seq_lens_lst):
165
block_table = [
166
- random.randint(0, NUM_BLOCKS - 1) for _ in range(max_num_blocks_per_seq)
166
+ random.randint(num_seqs, NUM_BLOCKS - 1)
167
+ for _ in range(max_num_blocks_per_seq)
168
]
169
+ # Keep the last block private to this sequence so its padding can be
170
+ # poisoned without corrupting valid tokens from another sequence.
171
+ block_table[(seq_len - 1) // block_size] = seq_idx
172
block_tables_lst.append(block_table)
173
174
block_tables = torch.tensor(block_tables_lst, dtype=torch.int)
@@ -183,8 +187,20 @@ def test_paged_attention(
187
)
188
key_cache, value_cache = key_caches[0], value_caches[0]
189
186
- # Using default kv_scale
187
- k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
190
+ k_scale_value = 0.5 if kv_cache_dtype == "fp8" else 1.0
191
+ v_scale_value = 0.25 if kv_cache_dtype == "fp8" else 1.0
192
+ k_scale = torch.tensor(k_scale_value, dtype=torch.float32, device=device)
193
+ v_scale = torch.tensor(v_scale_value, dtype=torch.float32, device=device)
194
+
195
+ if kv_cache_dtype == "fp8":
196
+ padding_nan = 0x80 if current_platform.is_fp8_fnuz() else 0x7F
197
+ else:
198
+ padding_nan = torch.nan
199
+ for seq_idx, seq_len in enumerate(seq_lens_lst):
200
+ padding_start = seq_len % block_size
201
+ if padding_start:
202
+ last_block_idx = block_tables_lst[seq_idx][seq_len // block_size]
203
+ value_cache[last_block_idx, :, :, padding_start:] = padding_nan
204
205
# Call the paged attention kernel.
206
output = torch.empty_like(query)
@@ -256,14 +272,14 @@ def test_paged_attention(
272
dequantized_key_cache = torch.empty(
273
size=key_cache_shape, dtype=dtype, device=device
274
)
259
- ops.convert_fp8(dequantized_key_cache, key_cache)
275
+ ops.convert_fp8(dequantized_key_cache, key_cache, k_scale_value)
276
key_cache = dequantized_key_cache
277
278
value_cache_shape = value_cache.shape
279
dequantized_value_cache = torch.empty(
280
size=value_cache_shape, dtype=dtype, device=device
281
)
266
- ops.convert_fp8(dequantized_value_cache, value_cache)
282
+ ops.convert_fp8(dequantized_value_cache, value_cache, v_scale_value)
283
value_cache = dequantized_value_cache
284
285
ref_output = torch.empty_like(query)