fix: use reshape instead of risky [0::2] slicing for E2M1 packing
This commit is contained in:
@@ -107,11 +107,10 @@ def _deepseek_v4_stage_mega_moe_inputs_kernel(
|
||||
|
||||
# Pack 2 E2M1 values per byte: even→low nibble, odd→high nibble
|
||||
PACKED_K: tl.constexpr = BLOCK_K // 2 # 64
|
||||
e2m1_4bit_flat = tl.reshape(e2m1_4bit, [BLOCK_K])
|
||||
# Interleave: index 0,2,4,... → low nibbles; 1,3,5,... → high nibbles
|
||||
even = e2m1_4bit_flat[0::2]
|
||||
odd = e2m1_4bit_flat[1::2]
|
||||
packed_byte = (odd.to(tl.uint8) << 4) | even.to(tl.uint8)
|
||||
e2m1_pairs = tl.reshape(e2m1_4bit, [PACKED_K, 2])
|
||||
even = e2m1_pairs[:, 0].to(tl.uint8)
|
||||
odd = e2m1_pairs[:, 1].to(tl.uint8)
|
||||
packed_byte = (odd << 4) | even
|
||||
|
||||
packed_k_offsets = k_block_id * PACKED_K + tl.arange(0, PACKED_K)
|
||||
packed_k_mask = packed_k_offsets < (hidden_size // 2)
|
||||
|
||||
Reference in New Issue
Block a user