[TRITON]: Standardize GEMM weight shape to (N, K) and TN memory layout (by default)

[willzhou-amd]

Changes:

• Fix .assume() bug that causes M=1 cases to fail in some kernels
• Add test cases for M = 1 to all tests
• Standardize weight shapes to (N, K)

[willzhou-amd]

There's a gnarly issue where Triton will implicitly promote small integers that it believes to be constants to tl.constant(). In some cases where the dimensions are small, typecasting arguments to your kernel will yield a cast error:

    batch_id = batch_id.to(tl.int64)
    stride_ab = stride_ab.to(tl.int64)
                ^
AttributeError("'constexpr' object has no attribute 'to'")

Replacing .to() with tl.cast() prevents this implicit promotion. I'm uncertain whether this happens on other versions of Triton, but this is reliably replicable on Triton 3.3.1.

tldr: If you're attempting to cast strides to int64 to prevent integer overflow, use tl.cast.