submission 183122

import torch
import triton
import triton.language as tl
import torch.nn.functional as F
from triton.tools.tensor_descriptor import TensorDescriptor
# Note : Selon ta version, l'import peut être 'triton.tools.tensor_descriptor'

# --- 1. KERNEL (Utilisation des Descriptors) ---

@triton.jit
def block_scaled_matmul_kernel(
    # Ici on reçoit les descripteurs (handles TMA)
    a_desc, b_desc, sfa_ptr, sfb_ptr, c_desc,
    # Dimensions (utiles pour les boucles et scales)
    M, N, K_packed,
    # Strides (uniquement pour les scales qui restent en pointeurs classiques)
    stride_sm, stride_sk,
    stride_sn, stride_skb,
    # Meta-paramètres
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    pid = tl.program_id(0)
    
    # --- Swizzling ---
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_am_desc = pid_m * BLOCK_M
    offs_bn_desc = pid_n * BLOCK_N
    

    # --- Offsets pour les Scales (Restent manuels) ---
    offs_k_scale = tl.arange(0, BLOCK_K // 16)
    offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    
    # Pointeurs Scales
    sfa_ptrs = sfa_ptr + offs_am[:, None] * stride_sm + offs_k_scale[None, :] * stride_sk
    sfb_ptrs = sfb_ptr + offs_bn[:, None] * stride_sn + offs_k_scale[None, :] * stride_skb

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # --- Boucle K ---
    # Avec TensorDescriptor, les offsets sont passés directement à la méthode .load()
    # K_packed est la dimension physique (bytes pour FP4)
    
    for k_idx in range(0, tl.cdiv(K_packed, BLOCK_K // 2)):
        # K offset en éléments packed (bytes)
        k_offset = k_idx * (BLOCK_K // 2)

        # --- LOAD VIA DESCRIPTOR (TMA) ---
        # L'objet descripteur gère le boundary check et le padding via le hardware (CuTensorMap)
        # On passe les coordonnées [Y, X] globales.
        
        a_tile = a_desc.load([offs_am_desc, k_offset])
        b_tile = b_desc.load([offs_bn_desc, k_offset])
        
        # Load scales (classique)
        sfa_tile = tl.load(sfa_ptrs)
        sfb_tile = tl.load(sfb_ptrs)

        # Compute
        accumulator = tl.dot_scaled(
            a_tile, sfa_tile, "e2m1",
            b_tile.T, sfb_tile, "e2m1",
            accumulator
        )

        # Advance scale pointers
        sfa_ptrs += (BLOCK_K // 16) * stride_sk
        sfb_ptrs += (BLOCK_K // 16) * stride_skb

    # --- STORE VIA DESCRIPTOR ---
    c_desc.store([offs_am_desc, offs_bn_desc], accumulator.to(tl.float16))


# --- 2. PYTHON LAUNCHER (Test Fixe) ---

def custom_kernel(data):
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data

    # 1. Vue uint8
    a = a.view(torch.uint8)
    b = b.view(torch.uint8)

    # 2. Nettoyage
    if a.dim() > 2: a = a.squeeze(-1)
    if b.dim() > 2: b = b.squeeze(-1)
    if sfa.dim() > 2: sfa = sfa.squeeze(-1)
    if sfb.dim() > 2: sfb = sfb.squeeze(-1)
    c_squeezed = c.squeeze(-1) if c.dim() > 2 else c

    M, K_packed = a.shape
    N, _ = b.shape

    # --- CONFIG FIXE ---
    BLOCK_M = 128
    BLOCK_N = 64
    BLOCK_K = 512
    GROUP_SIZE_M = 32
    
   

    M, K_packed_new = a.shape
    
    # 4. CRÉATION DES DESCRIPTEURS (CPU Side)
    # On instancie les TensorDescriptors ici avec les tailles de blocs fixes.
    # C'est ce que faisait ta lambda, mais on le fait en "dur" pour le test.
    
    a_desc = TensorDescriptor.from_tensor(a, [BLOCK_M, BLOCK_K // 2])
    b_desc = TensorDescriptor.from_tensor(b, [BLOCK_N, BLOCK_K // 2])
    c_desc = TensorDescriptor.from_tensor(c_squeezed, [BLOCK_M, BLOCK_N])

    # 5. Grid
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), )

    # 6. Lancement
    block_scaled_matmul_kernel[grid](
        # On passe les objets descripteurs directement
        a_desc, b_desc, sfa, sfb, c_desc,
        int(M), int(N), int(K_packed_new),
        sfa.stride(0), sfa.stride(1),
        sfb.stride(0), sfb.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        GROUP_SIZE_M=GROUP_SIZE_M,
        num_stages=4,
        num_warps=4
    )
    
    return c