diff --git a/Dockerfile b/Dockerfile index fe4962ff..8155124b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -39,7 +39,6 @@ ARG VLLM_MODELS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/ ARG VLLM_LAYERS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers COPY vllm/patches/deepseek_v4.py ${VLLM_MODELS_DIR}/deepseek_v4.py -COPY vllm/patches/staging_kernel.py ${VLLM_MODELS_DIR}/staging_kernel.py COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attention.py RUN sed -i 's/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),\n "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),/' \ diff --git a/src/nvfp4_megamoe_kernel/symm_buffer.py b/src/nvfp4_megamoe_kernel/symm_buffer.py index 84be7ae6..7c46658f 100644 --- a/src/nvfp4_megamoe_kernel/symm_buffer.py +++ b/src/nvfp4_megamoe_kernel/symm_buffer.py @@ -36,17 +36,17 @@ class SymmBuffer: device = torch.cuda.current_device() # NVFP4: packed E2M1 (2 values per byte), so K//2 - sf_k_groups_hidden = hidden_size // (16 * 4) # UE4M3 block16, 4 packed per uint32 - sf_k_groups_inter = intermediate_size // (16 * 4) + sf_k_groups_hidden = hidden_size // 16 # UE4M3 block16, 1 scale per group + sf_k_groups_inter = intermediate_size // 16 - # Staging buffers (matching DeepGEMM layout) + # Staging buffers (matching kernel's expected input format) self.x = torch.empty( max_num_tokens, hidden_size // 2, dtype=torch.int8, device=device, ) self.x_sf = torch.empty( max_num_tokens, sf_k_groups_hidden, - dtype=torch.uint32, device=device, + dtype=torch.float8_e4m3fn, device=device, ) self.topk_idx = torch.empty( max_num_tokens, top_k, diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 110dcb1f..8d50d383 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -206,9 +206,6 @@ class DeepseekV4FP8Config(Fp8Config): return isinstance(layer, FusedMoE) and self.expert_dtype == "fp4" -# Staging kernel imported from standalone module (has proper RNE rounding -# and subnormal handling for UE4M3 — the old embedded copy was broken). -from vllm.model_executor.models.staging_kernel import _stage_deepseek_v4_mega_moe_inputs def make_deepseek_v4_expert_params_mapping( @@ -532,15 +529,17 @@ class DeepseekV4MegaMoEExperts(nn.Module): symm_buffer = self.get_symm_buffer() symm_buffer.experts_start_idx = self.experts_start_idx num_tokens = hidden_states.shape[0] - _stage_deepseek_v4_mega_moe_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens], - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - ) + + # Quantize activation using the kernel's PyTorch stage_activation + # (same code path the kernel uses for L1→L2 requantization). + # This replaces the broken Triton staging kernel — no more uint32 + # pack/unpack, no more Triton tensor indexing issues. + from nvfp4_megamoe_kernel import stage_activation + x_fp4, x_sf = stage_activation(hidden_states) + symm_buffer.x[:num_tokens].copy_(x_fp4) + symm_buffer.x_sf[:num_tokens].copy_(x_sf) + symm_buffer.topk_idx[:num_tokens].copy_(topk_ids) + symm_buffer.topk_weights[:num_tokens].copy_(topk_weights) # Debug: check staging output import os diff --git a/vllm/patches/staging_kernel.py b/vllm/patches/staging_kernel.py deleted file mode 100644 index ad3d8358..00000000 --- a/vllm/patches/staging_kernel.py +++ /dev/null @@ -1,276 +0,0 @@ -""" -NVFP4 staging kernel — full FP4 (E2M1) activations + UE4M3 block16 scales. - -The mxf4nvf4 PTX instruction requires BOTH A and B to be FP4 (E2M1 packed). -This kernel quantizes BF16 activations → E2M1 packed uint8 with UE4M3 scales. -""" -import triton -import triton.language as tl -import torch - - -@triton.jit -def _deepseek_v4_stage_mega_moe_inputs_kernel( - hidden_states, - x_fp4, # uint8, shape (M, K//2) — E2M1 packed, 2 values per byte - x_sf, # int32, shape (M, K//64) — UE4M3 packed, 4 scales per int32 - topk_ids, - topk_weights, - topk_idx_out, - topk_weights_out, - hidden_stride_m: tl.constexpr, - hidden_stride_k: tl.constexpr, - x_stride_m: tl.constexpr, - x_stride_k: tl.constexpr, - x_sf_stride_m: tl.constexpr, - x_sf_stride_k: tl.constexpr, - topk_ids_stride_m: tl.constexpr, - topk_ids_stride_k: tl.constexpr, - topk_weights_stride_m: tl.constexpr, - topk_weights_stride_k: tl.constexpr, - topk_idx_stride_m: tl.constexpr, - topk_idx_stride_k: tl.constexpr, - topk_weights_out_stride_m: tl.constexpr, - topk_weights_out_stride_k: tl.constexpr, - hidden_size: tl.constexpr, - top_k: tl.constexpr, - BLOCK_K: tl.constexpr, # 128 elements (loaded from hidden) - GROUP_K: tl.constexpr, # 16 (NVFP4 group_size) - BLOCK_TOPK: tl.constexpr, -) -> None: - token_id = tl.program_id(0) - k_block_id = tl.program_id(1) - - k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) - k_mask = k_offsets < hidden_size - hidden = tl.load( - hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k, - mask=k_mask, - other=0.0, - ).to(tl.float32) - - num_groups: tl.constexpr = BLOCK_K // GROUP_K # 8 - hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K]) - abs_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K]) - amax = tl.max(abs_groups, axis=1) - amax = tl.maximum(amax, 1.0e-4) - - # ---- UE4M3 scale computation ---- - # scale = amax / 6.0 (E2M1 max value = 6) - # Then quantize scale to UE4M3 format - scale = amax / 6.0 - scale_bits = scale.to(tl.uint32, bitcast=True) - scale_exp = (scale_bits >> 23) & 0xFF - scale_mant = scale_bits & 0x7FFFFF - - # Convert FP32 → E4M3 manually (with subnormal support) - # FP32 bias=127, E4M3 bias=7 → raw exp = scale_exp - 120 - e4m3_exp_raw = scale_exp - 120 # can be negative → subnormal - - # Normal path: exp >= 1, just truncate mantissa top 3 bits - # RNE rounding: need guard (bit 19), sticky (OR of bits 18:0), and LSB of result - normal_mant = scale_mant >> 20 - guard_bit = (scale_mant >> 19) & 1 - sticky_bit = tl.where((scale_mant & 0x7FFFF) != 0, 1, 0) # OR of bits [18:0] - result_lsb = normal_mant & 1 - # RNE: round up if (guard=1 and sticky=1) or (guard=1 and sticky=0 and lsb=1) - round_up = guard_bit & (sticky_bit | result_lsb) - normal_mant = normal_mant + round_up - normal_exp = e4m3_exp_raw - - # Subnormal path: exp_raw <= 0 - # Insert implicit leading 1 and right-shift by (1 - exp_raw) - # E4M3 subnormal: value = (mant/8) * 2^(1-7) = (mant/8) * 2^-6 - # So we need: (1 + mant_fp32/2^23) * 2^(exp_raw - 7) = (shifted_mant/8) * 2^-6 - # shifted_mant = (implicit_1 | mant_fp32) >> (1 - exp_raw - 1) then take top 3 bits - shift = 1 - e4m3_exp_raw # positive when subnormal - mant_with_leading = (0x800000 | scale_mant) # insert implicit 1 - # Right-shift to get into the 3-bit E4M3 mantissa window - # We want bits [shift+19 : shift+23) of mant_with_leading for 3 mantissa bits + 1 round bit - subnormal_mant = (mant_with_leading >> (shift.to(tl.int32) + 20)) & 0x7 - sub_guard_bit = (mant_with_leading >> (shift.to(tl.int32) + 19)) & 1 - # Sticky: OR of all bits below the guard bit in the shifted result - # shift ≤ 8 in practice (amax floor = 1e-4 → scale ≈ 2^-15 → exp_raw ≈ -7), so mask ≤ 2^27 - sub_sticky_mask = (1 << (shift.to(tl.int32) + 19)) - 1 - sub_sticky_bit = tl.where((mant_with_leading & sub_sticky_mask) != 0, 1, 0) - sub_result_lsb = subnormal_mant & 1 - sub_round_up = sub_guard_bit & (sub_sticky_bit | sub_result_lsb) - subnormal_mant = subnormal_mant + sub_round_up - - is_normal = e4m3_exp_raw >= 1 - e4m3_mant = tl.where(is_normal, normal_mant, subnormal_mant) - e4m3_exp = tl.where(is_normal, normal_exp, 0) # exp=0 for subnormals - - # Handle mantissa overflow after rounding - overflow = e4m3_mant >= 8 - e4m3_mant = tl.where(overflow, 0, e4m3_mant) - e4m3_exp = tl.where(overflow, e4m3_exp + 1, e4m3_exp) - e4m3_exp = tl.maximum(e4m3_exp, 0) - # Saturation: E4M3FN reserves exp=15 for Inf/NaN (0x7F = NaN). - # Clamp to max representable finite value (exp=14, mant=7 = 0x77 = 448.0). - # This matches PyTorch's .to(torch.float8_e4m3fn) behavior. - sat = e4m3_exp >= 15 - e4m3_exp = tl.where(sat, 14, e4m3_exp) - e4m3_mant = tl.where(sat, 7, e4m3_mant) - scale_e4m3_bits = (e4m3_exp << 3) | e4m3_mant - - # Reconstruct dequantized scale by decoding the STORED E4M3 bits. - # This guarantees the E2M1 quantization divides by exactly the value - # the CUDA kernel will multiply back — same bits, single decode, no - # possibility of encode/decode disagreement. - stored_exp = (scale_e4m3_bits >> 3) & 0xF - stored_mant = scale_e4m3_bits & 0x7 - e4m3_exp_for_recon = tl.maximum(stored_exp.to(tl.int32) - 7, -126) - two_pow_exp_bits = (e4m3_exp_for_recon + 127).to(tl.uint32) << 23 - two_pow_exp = two_pow_exp_bits.to(tl.float32, bitcast=True) - normal_value = (1.0 + stored_mant.to(tl.float32) / 8.0) * two_pow_exp - subnormal_value = (stored_mant.to(tl.float32) / 8.0) * 0.015625 - e4m3_value = tl.where(stored_exp == 0, subnormal_value, normal_value) - - # ---- E2M1 FP4 quantization (unpacked, 1 byte/element) ---- - # E2M1 LUT (unsigned): [0, 0.5, 1, 1.5, 2, 3, 4, 6] - # Nearest-neighbor using thresholds (midpoints between consecutive values) - scaled = hidden_groups * (1.0 / tl.maximum(e4m3_value, 1e-6))[:, None] - # Clamp to E2M1 range [-6, 6] - scaled = tl.maximum(scaled, -6.0) - scaled = tl.minimum(scaled, 6.0) - - abs_s = tl.abs(scaled) - # Thresholds: midpoints between [0, 0.5, 1, 1.5, 2, 3, 4, 6] - # [0, 0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, INF] - e2m1_idx = tl.where(abs_s < 0.25, 0, - tl.where(abs_s < 0.75, 1, - tl.where(abs_s < 1.25, 2, - tl.where(abs_s < 1.75, 3, - tl.where(abs_s < 2.5, 4, - tl.where(abs_s < 3.5, 5, - tl.where(abs_s < 5.0, 6, 7))))))) - sign_bit = (scaled < 0).to(tl.int32) - e2m1_4bit = (sign_bit << 3) | e2m1_idx # 4-bit: (sign << 3) | index - - # Pack E2M1 pairs into single bytes (2 per byte, low nibble first) - # mxf4nvf4 reads FP4 packed from SMEM — must match kernel's TMA layout - # e2m1_4bit is [num_groups, GROUP_K] — stride within each group (row-major - # layout means within-group striding pairs the same elements as flat striding) - e2m1_lo = e2m1_4bit[:, 0::2] # even within group → low nibble - e2m1_hi = e2m1_4bit[:, 1::2] # odd within group → high nibble - e2m1_packed = (e2m1_hi << 4 | e2m1_lo).to(tl.uint8) # [BLOCK_K // 2] - - k_offsets_out = k_block_id * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2) - k_mask_out = k_offsets_out < (hidden_size // 2) - tl.store( - x_fp4 + token_id * x_stride_m + k_offsets_out * x_stride_k, - e2m1_packed, - mask=k_mask_out, - ) - - # Pack UE4M3 bytes into int32 (NVFP4: group_size=16, 4 groups per 64 elements) - # 8 groups per k_block of 128 → 2 int32s per k_block - # int32 can only pack 4 bytes (shifts >= 32 are UB), so split into two packs - scale_offsets = tl.arange(0, num_groups) # [0..7] - first_half = scale_offsets < 4 # groups 0-3 → int32[0] - second_half = scale_offsets >= 4 # groups 4-7 → int32[1] - - packed_lo = tl.sum( - tl.where(first_half, scale_e4m3_bits.to(tl.int32) << (scale_offsets * 8), 0), - axis=0, - ).to(tl.int32) - packed_hi = tl.sum( - tl.where(second_half, scale_e4m3_bits.to(tl.int32) << ((scale_offsets - 4) * 8), 0), - axis=0, - ).to(tl.int32) - - # Write 2 int32s per k_block: x_sf shape is (M, K//64) = (M, num_k_blocks * 2) - sf_base = token_id * x_sf_stride_m + k_block_id * 2 * x_sf_stride_k - tl.store(x_sf + sf_base, packed_lo) - tl.store(x_sf + sf_base + x_sf_stride_k, packed_hi) - - if k_block_id == 0: - topk_offsets = tl.arange(0, BLOCK_TOPK) - topk_mask = topk_offsets < top_k - - ids = tl.load( - topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k, - mask=topk_mask, - other=0, - ).to(tl.int32) - tl.store( - topk_idx_out - + token_id * topk_idx_stride_m - + topk_offsets * topk_idx_stride_k, - ids, - mask=topk_mask, - ) - - weights = tl.load( - topk_weights - + token_id * topk_weights_stride_m - + topk_offsets * topk_weights_stride_k, - mask=topk_mask, - other=0.0, - ) - tl.store( - topk_weights_out - + token_id * topk_weights_out_stride_m - + topk_offsets * topk_weights_out_stride_k, - weights, - mask=topk_mask, - ) - - -def _stage_deepseek_v4_mega_moe_inputs( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - x_fp4: torch.Tensor, # uint8, shape (M, K//2) - x_sf: torch.Tensor, # int32, shape (M, K//64) - topk_idx_out: torch.Tensor, - topk_weights_out: torch.Tensor, -) -> None: - num_tokens, hidden_size = hidden_states.shape - if num_tokens == 0: - return - if hidden_size % 128 != 0: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires hidden_size to be " - "a multiple of 128." - ) - top_k = topk_ids.shape[1] - if topk_weights.shape != topk_ids.shape: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires topk_weights and " - "topk_ids to have the same shape." - ) - - block_k = 128 - grid = (num_tokens, triton.cdiv(hidden_size, block_k)) - block_topk = triton.next_power_of_2(top_k) - _deepseek_v4_stage_mega_moe_inputs_kernel[grid]( - hidden_states, - x_fp4, - x_sf, - topk_ids, - topk_weights, - topk_idx_out, - topk_weights_out, - hidden_states.stride(0), - hidden_states.stride(1), - x_fp4.stride(0), - x_fp4.stride(1), - x_sf.stride(0), - x_sf.stride(1), - topk_ids.stride(0), - topk_ids.stride(1), - topk_weights.stride(0), - topk_weights.stride(1), - topk_idx_out.stride(0), - topk_idx_out.stride(1), - topk_weights_out.stride(0), - topk_weights_out.stride(1), - hidden_size, - top_k, - BLOCK_K=block_k, - GROUP_K=16, # NVFP4: group_size=16 (scale_vec::4X) - BLOCK_TOPK=block_topk, - num_warps=4, - )