refactor: accept arbitrary K in bf16 gemv with aligned head-tail sweeps
- Drop the K % 2 entry rejection and the per-K if/else load-width branch: the weight stream now anchors uint4 loads at each row's first 16-byte-aligned address, with scalar head/tail sweeps covering at most 14 remainder elements, so any positive K and any storage offset is correct - Keep one pure-uint4 loop (no branching inside the loop) for the production case where every x row base is 16-byte aligned (K % 8 == 0 with allocator-aligned tensors) and a scalar-x pairing loop only for unaligned K, where per-row uint4 loads are not addressable; measured cost of scalar x everywhere was up to 2.5x on multi-row shapes (down M=4 28.4us vs 11.3us) - Remove the now-obsolete k_aligned axis and K divisibility gate from the linear dispatch spec since the primitive no longer rejects any K - Add test coverage for unaligned K (7, 12, 100, 1534) at M=1 and M=3 Benchmark: 8x L20 (sm_89, CUDA 12.8), L2-resident microbench, 300 iters; hot path unchanged within noise vs the pure-uint4 kernel (q M=2 5.8us, down M=4 11.3us, lm M=1 391us); full gate green
This commit is contained in:
@@ -107,7 +107,6 @@ def _axes(
|
||||
x_contiguous=x.is_contiguous(),
|
||||
weight_contiguous=weight.is_contiguous(),
|
||||
bias_supported=bias_supported,
|
||||
k_even=k is not None and k % 2 == 0,
|
||||
)
|
||||
|
||||
|
||||
@@ -122,7 +121,6 @@ _SPEC_CAPABLE = (
|
||||
& axis("x_contiguous").truthy()
|
||||
& axis("weight_contiguous").truthy()
|
||||
& axis("bias_supported").truthy()
|
||||
& axis("k_even").truthy()
|
||||
)
|
||||
|
||||
_SPEC_AUTO = _SPEC_CAPABLE & Spec.of(
|
||||
@@ -166,7 +164,6 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
|
||||
or x.ndim not in (1, 2)
|
||||
or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
|
||||
or x.shape[-1] != weight.shape[1]
|
||||
or weight.shape[1] % 2 != 0
|
||||
or x.device != weight.device
|
||||
or not x.is_contiguous()
|
||||
or not weight.is_contiguous()
|
||||
|
||||
Reference in New Issue
Block a user