From 076d325c97c33fc578b7f9a7f0d7144ec7ef3297 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 11 May 2026 21:04:53 +0000 Subject: [PATCH] fix: use reshape instead of risky [0::2] slicing for E2M1 packing --- patches/staging_kernel.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/patches/staging_kernel.py b/patches/staging_kernel.py index 83476f3..7d18681 100644 --- a/patches/staging_kernel.py +++ b/patches/staging_kernel.py @@ -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)