fix: restore l parameter in wrapper for backward compat when b_tensor_l_sizes is None

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: 2 people

flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py

@@ -410,8 +410,8 @@ def __init__(
         vectorized_f32: bool,
         topk: cutlass.Int64,
         raster_along_m: bool = False,
-        enable_pdl: bool = True,
         b_tensor_l_sizes: Optional[Tuple[int, ...]] = None,
+        enable_pdl: bool = True,
     ):
         """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel with
         gather operation and SwiGLU fusion.
@@ -533,25 +533,24 @@ def __init__(
         self.vectorized_f32 = vectorized_f32
 
         # Multi-B tensor configuration
-        # b_tensor_l_sizes is required — the Python wrapper layer always provides it
-        # as a tuple (even for single-B, e.g. (256,)).
         if b_tensor_l_sizes is None:
-            raise ValueError(
-                "b_tensor_l_sizes is required. Pass a tuple with the number of "
-                "experts per tensor, e.g. (num_experts,) for single-B."
+            self.num_b_tensors = 1
+            self.b_tensor_l_sizes = None
+            # Offsets padded for safe indexing in kernel
+            self.b_tensor_l_offsets = (0,) + (2**30,) * self.MAX_B_TENSORS
+        else:
+            assert len(b_tensor_l_sizes) <= self.MAX_B_TENSORS, (
+                f"Max {self.MAX_B_TENSORS} B tensors, got {len(b_tensor_l_sizes)}"
             )
-        assert len(b_tensor_l_sizes) <= self.MAX_B_TENSORS, (
-            f"Max {self.MAX_B_TENSORS} B tensors, got {len(b_tensor_l_sizes)}"
-        )
-        self.num_b_tensors = len(b_tensor_l_sizes)
-        self.b_tensor_l_sizes = b_tensor_l_sizes
-        offsets = [0]
-        for l_size in b_tensor_l_sizes:
-            offsets.append(offsets[-1] + l_size)
-        # Pad to MAX_B_TENSORS + 1 for safe indexing
-        while len(offsets) < self.MAX_B_TENSORS + 1:
-            offsets.append(2**30)
-        self.b_tensor_l_offsets = tuple(offsets)
+            self.num_b_tensors = len(b_tensor_l_sizes)
+            self.b_tensor_l_sizes = b_tensor_l_sizes
+            offsets = [0]
+            for l_size in b_tensor_l_sizes:
+                offsets.append(offsets[-1] + l_size)
+            # Pad to MAX_B_TENSORS + 1 for safe indexing
+            while len(offsets) < self.MAX_B_TENSORS + 1:
+                offsets.append(2**30)
+            self.b_tensor_l_offsets = tuple(offsets)
 
     def _setup_attributes(self):
         """Set up configurations that are dependent on GEMM inputs
@@ -4034,6 +4033,7 @@ def wrapper(
         m: cutlass.Int64,
         n: cutlass.Int64,
         k: cutlass.Int64,
+        l: cutlass.Int64,  # noqa: E741
         tile_size: cutlass.Constexpr,
         scaling_vector_size: cutlass.Constexpr,
         max_active_clusters: cutlass.Constexpr,
@@ -4043,12 +4043,19 @@ def wrapper(
         """Unified wrapper supporting both single-B and multi-B tensors.
 
         B tensors are always passed as tuples (length 1 for single-B).
-        L sizes are configured via b_tensor_l_sizes in __init__.
+        When b_tensor_l_sizes is provided, L sizes come from b_tensor_l_sizes;
+        otherwise falls back to the l parameter (backward compatible single-B).
         """
         scale_k = k // scaling_vector_size
         interm_size = n // 2
         num_tiles = m // tile_size
-        total_l = self.b_tensor_l_offsets[self.num_b_tensors]
+        # When b_tensor_l_sizes is provided, total_l comes from the precomputed offsets
+        # and l is ignored. Callers must ensure l == sum(b_tensor_l_sizes).
+        # When b_tensor_l_sizes is None (single-B backward compat), l is used directly.
+        if cutlass.const_expr(self.b_tensor_l_sizes is not None):
+            total_l = self.b_tensor_l_offsets[self.num_b_tensors]
+        else:
+            total_l = l
 
         a = cute.make_tensor(
             a_ptr, layout=cute.make_ordered_layout((orig_m, k, 1), order=(1, 0, 2))
@@ -4069,7 +4076,10 @@ def wrapper(
         )
 
         # Create B and alpha tensors using const_expr conditions
-        l_0 = self.b_tensor_l_sizes[0]
+        if cutlass.const_expr(self.b_tensor_l_sizes is not None):
+            l_0 = self.b_tensor_l_sizes[0]
+        else:
+            l_0 = l
         alpha_0 = cute.make_tensor(alpha_ptr_tuple[0], layout=cute.make_layout((l_0,)))
         b_0 = cute.make_tensor(
             b_ptr_tuple[0],

flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py

@@ -274,8 +274,9 @@ def _get_compiled_gather_kernel(
         # Order must match wrapper signature:
         # (a_ptr, b_ptr_tuple, a_sf_ptr, b_sf_ptr_tuple, c_ptr, c_sf_ptr, alpha_ptr_tuple,
         #  tile_idx_to_group_idx_ptr, tile_idx_to_mn_limit_ptr, token_id_mapping_ptr,
-        #  num_non_exiting_tiles_ptr, norm_const_ptr, orig_m, m, n, k,
+        #  num_non_exiting_tiles_ptr, norm_const_ptr, orig_m, m, n, k, l,
         #  tile_size, scaling_vector_size, max_active_clusters, stream)
+        num_experts = sum(b_tensor_l_sizes)
         compile_args = [
             a_ptr,
             b_ptr,
@@ -293,6 +294,7 @@ def _get_compiled_gather_kernel(
             permuted_m,
             n,
             k,
+            num_experts,
         ]
 
         compiled_gemm = cute.compile(
@@ -620,6 +622,7 @@ def blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion_nvfp4(
     )
 
     # Execute kernel
+    num_experts = sum(b_tensor_l_sizes)
     exec_args = [
         a_ptr,
         b_ptr,
@@ -637,6 +640,7 @@ def blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion_nvfp4(
         permuted_m,
         n,
         k,
+        num_experts,  # l
     ]
     compiled_gemm(*exec_args, stream=stream)