refactor: replace magic layout ints with TensorLayout enum
- Add TensorLayout enum (C++ + Python) to replace magic layout ints - Add C10_CUDA_CHECK post-launch error checking to all kernel entries - Add CUDAGuard + freqs_cis shape validation to rotary_emb.cu - Cache SM count to eliminate per-call cudaDeviceGetAttribute - Add DISPATCH_CAUSAL_MASK macro to deduplicate dispatcher if/else - Convert mask type hints from X|None to Optional[X]
This commit is contained in:
@@ -1,5 +1,13 @@
|
||||
#pragma once
|
||||
|
||||
// Tensor layout for Q/K/V tensors passed to attention kernels.
|
||||
// Internally, kernels always operate on BHLD [batch, n_heads, seq_len, head_dim].
|
||||
// When the caller passes BLHD, dims 1 and 2 are transposed at entry.
|
||||
enum TensorLayout : int {
|
||||
BHLD = 0, // [batch, n_heads, seq_len, head_dim]
|
||||
BLHD = 1, // [batch, seq_len, n_heads, head_dim]
|
||||
};
|
||||
|
||||
|
||||
template<typename T, typename AT = float>
|
||||
struct AttentionParams {
|
||||
|
||||
Reference in New Issue
Block a user