Merge pull request #20 from lfr-0531/user/fanrongl/fix_scaling_factor

Opt padding+quant kernel

Author: Barry-Delaney

tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py

@@ -25,53 +25,55 @@ def _masked_index_copy_group_quant_fp8(
     # mask indices
     start_offsets_ptr,
     row_indices_ptr,
-    # group size
-    group_size,
-    # output size
+    # dimensions
+    num_groups,
     row_size,
     col_size,
     dim_size,
-    # avoid to divide zero
+    group_size,
+    # quantization parameters
     eps,
+    fp8_max,
     # block size
     BLOCK: tl.constexpr,
+    NUM_STAGE: tl.constexpr,
 ):
-    # get program id and block offset
-    pid = tl.program_id(0)
-    block_start = pid * group_size
+    group_block = tl.program_id(0)
+    token_block = tl.program_id(1)
+    token_block_num = tl.num_programs(1)
 
-    # compute mask and pointers
-    offsets = block_start + tl.arange(0, BLOCK)
-    mask = offsets < (block_start + group_size)
+    # calculate group and element offsets
     num_tokens = tl.load(start_offsets_ptr + row_size)
-    token_idx = offsets // dim_size
-    valid = (token_idx < num_tokens) & mask
-    row_idx = tl.load(row_indices_ptr + token_idx, mask=valid)
-    start_offset = tl.load(start_offsets_ptr + row_idx, mask=valid)
-    col_idx = token_idx - start_offset
-    elem_idx = offsets % dim_size
-
-    # load input data
-    input = tl.load(input_ptr + offsets, mask=valid, other=0.0).to(tl.float32)
-
-    # quant
-    _absmax = tl.maximum(tl.max(tl.abs(input)), eps)
-    output_s = _absmax / 448.0
-    output_s = tl.exp2(tl.ceil(tl.log2(tl.abs(output_s))))
-    output_s_inv = 1.0 / output_s
-    output_q = tl.clamp(input * output_s_inv, -448.0,
-                        448.0).to(out_q_ptr.dtype.element_ty)
-
-    # write output
-    s_dim_size = dim_size // group_size
-    out_offsets = row_idx * col_size * dim_size + col_idx * dim_size + elem_idx
-    group_in_token = elem_idx // group_size
-    out_s_offset = row_idx * col_size * s_dim_size + col_idx * s_dim_size + group_in_token
-
-    # Only store scaling factor for the first element in each group to avoid race conditions
-    is_first_in_group = elem_idx % group_size == 0
-    tl.store(out_q_ptr + out_offsets, output_q, mask=valid)
-    tl.store(out_s_ptr + out_s_offset, output_s, mask=valid & is_first_in_group)
+    group_start = group_block * group_size
+    elem_offsets = group_start + tl.arange(0, BLOCK)
+    valid_elem = elem_offsets < (group_start + group_size)
+    input_ptr_offs = input_ptr + elem_offsets
+    output_ptr_offs = out_q_ptr + elem_offsets
+    output_s_offs = out_s_ptr + group_block
+
+    # process tokens
+    for token_index in tl.range(token_block,
+                                num_tokens,
+                                token_block_num,
+                                num_stages=NUM_STAGE):
+        # load input and indices
+        input_data = tl.load(input_ptr_offs + token_index * dim_size,
+                             mask=valid_elem,
+                             other=0.0)
+        row_idx = tl.load(row_indices_ptr + token_index)
+        start_offset = tl.load(start_offsets_ptr + row_idx)
+        idx = row_idx * col_size + token_index - start_offset
+
+        # quantization
+        _absmax = tl.maximum(tl.max(tl.abs(input_data)), eps)
+        output_s = _absmax / fp8_max
+        output_s = tl.exp2(tl.ceil(tl.log2(tl.abs(output_s))))
+        output_q = tl.clamp(input_data / output_s, -fp8_max,
+                            fp8_max).to(out_q_ptr.dtype.element_ty)
+
+        # store quantized values and scaling factor
+        tl.store(output_ptr_offs + idx * dim_size, output_q, mask=valid_elem)
+        tl.store(output_s_offs + idx * num_groups, output_s)
 
 
 def masked_index_copy_group_quant_fp8(
@@ -88,32 +90,50 @@ def masked_index_copy_group_quant_fp8(
     ), "the last dimension of `input` cannot be divisible by `group_size`"
     assert input.is_contiguous(), "`input` is not contiguous"
     assert input.ndim == 2, "Input must be a 2D tensor"
-    assert output.ndim == 3, "Input must be a 3D tensor, [row, col, dim]"
+    assert output.ndim == 3, "Output must be a 3D tensor, [row, col, dim]"
     assert start_offsets.shape[
         0] == output.shape[0] + 1, "Start offsets must be (num_experts + 1)"
 
     num_tokens = input.shape[0]
     row_size = output.shape[0]
     col_size = output.shape[1]
     dim_size = output.shape[2]
-    total_elems = num_tokens * dim_size
+    num_groups = (dim_size + group_size - 1) // group_size
+
+    # get block/grid/stage/warp
+    BLOCK = group_size
+    if num_tokens <= 4096:
+        TOKEN_BLOCK_NUM = 128
+        NUM_STAGES = 4
+        num_warps = 2
+    else:
+        TOKEN_BLOCK_NUM = 64
+        NUM_STAGES = 6
+        num_warps = 1
+    grid = (
+        num_groups,
+        TOKEN_BLOCK_NUM,
+    )
 
-    M = total_elems // group_size
-    BLOCK = triton.next_power_of_2(group_size)
-    # heuristics for number of warps
-    num_warps = min(max(BLOCK // 256, 1), 8)
-    _masked_index_copy_group_quant_fp8[(M, )](
+    # FP8 quantization parameters
+    finfo = torch.finfo(torch.float8_e4m3fn)
+    fp8_max = finfo.max
+
+    _masked_index_copy_group_quant_fp8[grid](
         input,
         output,
         output_s,
         start_offsets,
         row_indices,
-        group_size,
+        num_groups,
         row_size,
         col_size,
         dim_size,
+        group_size,
         eps,
+        fp8_max,
         BLOCK=BLOCK,
+        NUM_STAGE=NUM_STAGES,
         num_warps=num_warps,
     )
     return