diff --git a/patches/deepseek_v4.py b/patches/deepseek_v4.py index 5958d08..9a51190 100644 --- a/patches/deepseek_v4.py +++ b/patches/deepseek_v4.py @@ -353,9 +353,8 @@ 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_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 + even, odd = tl.split(e2m1_pairs) # splits last axis (size 2) into two [PACKED_K] tensors + packed_byte = (odd.to(tl.uint8) << 4) | even.to(tl.uint8) packed_k_offsets = k_block_id * PACKED_K + tl.arange(0, PACKED_K) packed_k_mask = packed_k_offsets < (hidden_size // 2)