[WIP] Before refactor the compute)_chunked_prefill.

tests/test_xattn_kernels.py

@@ -6,9 +6,10 @@ Test: XAttention Triton kernels
 2. softmax_fuse_block_sum: 对 attention scores 做 softmax 后按 block 求和
 
 数据流:
-  Q, K [batch, heads, seq_len, head_dim]
+  Q [batch, heads, q_len, head_dim]
+  K [batch, heads, kv_len, head_dim]
     ↓ flat_group_gemm_fuse_reshape
-  attn_scores [batch, heads, seq_len/stride, seq_len/stride]
+  attn_scores [batch, heads, q_len/stride, kv_len/stride]
     ↓ softmax_fuse_block_sum
   block_sums [batch, heads, q_blocks, k_blocks]
 """
@@ -21,7 +22,11 @@ from nanovllm.ops.xattn import flat_group_gemm_fuse_reshape, softmax_fuse_block_
 # 参数配置
 # ============================================================
 
-seq_len = 512       # Triton 要求 seq_len >= stride * BLOCK_M = 4 * 128 = 512
+# Triton 约束: q_len >= stride * BLOCK_M, kv_len >= stride * BLOCK_N
+# A100: BLOCK_M = BLOCK_N = 128, 所以 min = 4 * 128 = 512
+# RTX 3090: BLOCK_M = BLOCK_N = 64, 所以 min = 4 * 64 = 256
+q_len = 512
+kv_len = 2048
 head_dim = 128
 stride = 4
 block_size = 128    # softmax block size (in reshaped space)
@@ -31,26 +36,56 @@ segment_size = 128  # Triton kernel 要求 segment_size >= block_size
 # 构造输入: 偶数位置=1, 奇数位置=2
 # ============================================================
 
-Q = torch.zeros(1, 1, seq_len, head_dim, dtype=torch.bfloat16).cuda()
-K = torch.zeros(1, 1, seq_len, head_dim, dtype=torch.bfloat16).cuda()
-for i in range(seq_len):
+Q = torch.zeros(1, 1, q_len, head_dim, dtype=torch.bfloat16).cuda()
+K = torch.zeros(1, 1, kv_len, head_dim, dtype=torch.bfloat16).cuda()
+
+for i in range(q_len):
     if i % 2 == 0:
         Q[0, 0, i, :] = 1
-        K[0, 0, i, :] = 1
     else:
         Q[0, 0, i, :] = 2
+
+for i in range(kv_len):
+    if i % 2 == 0:
+        K[0, 0, i, :] = 1
+    else:
         K[0, 0, i, :] = 2
 
 # ============================================================
-# Step 1: flat_group_gemm_fuse_reshape
+# Step 1: flat_group_gemm_fuse_reshape (chunked along K)
 # ============================================================
 
-attn_scores = flat_group_gemm_fuse_reshape(
-    Q, K, stride,
-    chunk_start=0,
-    chunk_end=seq_len // stride,
-    is_causal=False
-)
+q_reshaped_len = q_len // stride   # 128
+kv_reshaped_len = kv_len // stride  # 512
+
+# 将 K 沿着长度维度分成多个 chunk
+k_chunk_size = 512  # 每个 chunk 512 tokens
+num_k_chunks = kv_len // k_chunk_size  # 4 chunks
+
+attn_scores_list = []
+for k_chunk_idx in range(num_k_chunks):
+    k_start = k_chunk_idx * k_chunk_size
+    k_end = k_start + k_chunk_size
+    K_chunk = K[:, :, k_start:k_end, :]  # [1, 1, k_chunk_size, head_dim]
+
+    # 对每个 K chunk 调用 flat_group_gemm_fuse_reshape
+    # 输出: [batch, heads, q_len/stride, k_chunk_size/stride]
+    attn_chunk = flat_group_gemm_fuse_reshape(
+        Q, K_chunk, stride,
+        chunk_start=0,
+        chunk_end=q_reshaped_len,
+        is_causal=False
+    )
+    attn_scores_list.append(attn_chunk)
+
+# 拼接所有 K chunks 的结果
+# 每个 chunk: [1, 1, q_reshaped_len, k_chunk_size/stride]
+# 拼接后: [1, 1, q_reshaped_len, kv_reshaped_len]
+attn_scores = torch.cat(attn_scores_list, dim=-1)
+
+# 验证 shape: [batch, heads, q_len/stride, kv_len/stride]
+assert attn_scores.shape == (1, 1, q_reshaped_len, kv_reshaped_len), \
+    f"shape mismatch: {attn_scores.shape} != (1, 1, {q_reshaped_len}, {kv_reshaped_len})"
 
 # 验证: 反对角线求和
 # 每个 stride x stride 块的反对角线: Q[奇]*K[偶] + Q[偶]*K[奇] = 2*1 + 1*2 = 4
@@ -63,7 +98,6 @@ assert actual_gemm == expected_gemm, f"flat_group_gemm: {actual_gemm} != {expect
 # Step 2: softmax_fuse_block_sum
 # ============================================================
 
-reshaped_len = seq_len // stride
 scale = 1.4426950408889634  # log2(e) for exp2
 
 block_sums = softmax_fuse_block_sum(
@@ -71,15 +105,24 @@ block_sums = softmax_fuse_block_sum(
     block_size,
     segment_size,
     chunk_start=0,
-    chunk_end=reshaped_len,
-    real_q_len=reshaped_len,
+    chunk_end=q_reshaped_len,
+    real_q_len=q_reshaped_len,
     scale=scale,
     is_causal=False
 )
 
+# 验证 shape: [batch, heads, q_blocks, k_blocks]
+q_blocks = q_reshaped_len // block_size  # 128 / 128 = 1
+k_blocks = kv_reshaped_len // block_size  # 512 / 128 = 4
+assert block_sums.shape == (1, 1, q_blocks, k_blocks), \
+    f"shape mismatch: {block_sums.shape} != (1, 1, {q_blocks}, {k_blocks})"
+
 # 验证: 每个 block 的 softmax 结果求和
-# 所有 attn_scores 相同 → softmax 均匀分布 → block_sum = block_size^2 / reshaped_len
-expected_sum = block_size * block_size / reshaped_len
+# 所有 attn_scores 相同 → softmax 均匀分布
+# 每行对一个 K block 的贡献 = block_size / kv_reshaped_len
+# 每个 Q block 有 block_size 行
+# block_sum = block_size * (block_size / kv_reshaped_len)
+expected_sum = block_size * block_size / kv_reshaped_len
 actual_sum = block_sums[0, 0, 0, 0].item()
 assert actual_sum == expected_sum, f"softmax_fuse_block_sum: {actual_sum} != {expected_sum}"