[Kernel] Keep SM70 attention probabilities in registers
Preserve FP32 accumulation and the original K64 online reduction order while eliminating shared probability traffic and cross-warp softmax synchronization. Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
This commit is contained in:
yangzhuxinyzx committed
1 parent
f62eff373a
commit
f8b85c681a
3 files changed
+172
-100
No files matched your search
@@ -0,0 +1,76 @@
|
||||
# FlashInfer SM70 probability fragments in registers
|
||||
|
||||
This change improves the explicitly selected FlashInfer backend. It does not
|
||||
qualify an AUTO configuration or meet the campaign's >80 useful TFLOP/s/card
|
||||
gate. The shared `sm70_attention.noncausal_attention` interface dispatches on
|
||||
SM70 hardware, FP16 tensors and BSHD D128 layout, independently of model,
|
||||
quantization and adapter identity.
|
||||
|
||||
## Arithmetic and execution
|
||||
|
||||
One warp owns 16 query rows and both logical 32-key halves of the existing
|
||||
64-key tile. QK/PV accumulators, online maxima and denominators remain FP32.
|
||||
Each original partial sum, XOR-2/XOR-8 reduction, left/right addition and
|
||||
per-output MMA K order is retained. Probabilities cross the same FP16 rounding
|
||||
boundary as before.
|
||||
|
||||
Six 32-bit lane exchanges convert the rounded probability pairs into Volta A
|
||||
fragments. PV reuses each fragment across eight output fragments. Probabilities
|
||||
no longer traverse shared memory, and the warp owns its softmax state without
|
||||
cross-warp barriers. CTA barriers still protect K/V staging and reuse.
|
||||
|
||||
The CTA covers 192 queries with 384 threads. Cooperative prefetch rounds up
|
||||
to three vectors per thread. Extra complete warps in the final prefetch group
|
||||
skip rows outside the 64-key tile; the guard is warp-uniform before shuffles.
|
||||
Valid-length padding, batches, heads and unaligned global-storage loads retain
|
||||
their original handling. Explicit lane selects avoid addressable local arrays.
|
||||
|
||||
CUDA 12.8 emits 168 registers/thread, no stack or spills, and 84,992 dynamic
|
||||
shared bytes. These resource counts describe the implementation; they are not
|
||||
throughput measurements.
|
||||
|
||||
The FP16-accumulator layout shortcut in [FastAttention Appendix B](https://arxiv.org/html/2410.16663v1#A2)
|
||||
is outside this campaign's precision contract. This implementation exchanges
|
||||
already rounded probabilities while retaining FP32 accumulation.
|
||||
|
||||
## Development evidence
|
||||
|
||||
Environment: Python 3.12.13, Torch 2.10.0+cu128, CUDA toolkit 12.8.93,
|
||||
V100 SXM2 32GB. Artifacts are under
|
||||
`/data/minimax-h3/sm70-general-20260909/`.
|
||||
|
||||
- `attention-fi-register-probability-q192/probe.json`: 17 boundary lengths
|
||||
and the actual 34,551-token, 14-head H3 capture match the frozen FI binary
|
||||
bitwise. Sampled FP32-reference relative L2 is 0.000374 on the actual input.
|
||||
- Seven alternating operator timings on the same GPU give medians
|
||||
183.124985 -> 160.701447 ms, 12.245% lower latency. Observed clocks remain
|
||||
1425-1432 MHz. This is an operator measurement only.
|
||||
- Compute Sanitizer 12.8 memcheck, racecheck and synccheck each report zero
|
||||
errors on 12 boundary cases with two batches, three heads and storage offsets.
|
||||
- The formatted native build passes 33 GPU numerical, unaligned-storage,
|
||||
independent-query, graph and non-H3 shared-interface checks. The affected
|
||||
CPU contracts and provenance/acceptance suite pass 46 checks, with seven
|
||||
GPU checks skipped in the explicitly device-masked CPU invocation.
|
||||
- `fi-general-control-quality.json`: the current common prepared/residual
|
||||
path preserves the previously frozen FI final video/audio latents bitwise.
|
||||
- `fi-register-denoise-summary.json`: one complete denoise warmup per
|
||||
implementation, followed by one unprofiled measurement per implementation
|
||||
in reverse order. Both use the same TP4 W8A16 LightX2V four-step v1.2 weights,
|
||||
five sigma points, captured conditioning/noise, column layout, exact residual
|
||||
sharding and zero persistent FP16 weight cache. All four final-latent pairs
|
||||
match bitwise on every rank.
|
||||
|
||||
| Complete-denoise control | Slowest-rank seconds | Useful TFLOP/s/card |
|
||||
| --- | ---: | ---: |
|
||||
| Previous FI kernel, common execution path | 67.303216 | 46.108983-46.109000 |
|
||||
| Register-probability FI kernel, same path | 62.471266 | 49.675363-49.675382 |
|
||||
|
||||
The isolated kernel change reduces this complete denoise by 7.179%. Candidate
|
||||
steps take about 15.59-15.64 seconds. The control excludes encoder, VAE and
|
||||
packaging, and includes only one measurement per implementation. It is not
|
||||
the required full-request warmup-plus-three acceptance. Existing FA query-128
|
||||
remains faster in its separately recorded full-request measurements.
|
||||
|
||||
Full native media preservation, formal repeated request measurements,
|
||||
independent official references and human review remain separate gates. The
|
||||
campaign remains incomplete. No precision gate or sampler setting is relaxed.
|
||||
@@ -9,28 +9,20 @@
|
||||
namespace fi = flashinfer::attention::sm70;
|
||||
namespace {
|
||||
constexpr int D = 128;
|
||||
constexpr int BQ = 128;
|
||||
constexpr int BQ = 192;
|
||||
constexpr int BK = 64;
|
||||
// Two warps share 16 query rows. Each owns two K16 score fragments, so
|
||||
// expanding BK does not double the CTA's thread count.
|
||||
constexpr int KEY_WARPS = 2;
|
||||
// One warp owns Q16 and both logical K32 halves, retaining FP32 arithmetic.
|
||||
constexpr int KEY_WARPS = 1;
|
||||
constexpr int KEY_FRAGMENTS = BK / (KEY_WARPS * 16);
|
||||
constexpr int OUTPUT_FRAGMENTS = D / (KEY_WARPS * 16);
|
||||
constexpr int VLD = BK + 8;
|
||||
constexpr int THREADS = (BQ / 16) * KEY_WARPS * 32;
|
||||
constexpr int PREFETCH_VECTORS = BK * D / (THREADS * 8);
|
||||
static_assert(THREADS * PREFETCH_VECTORS * 8 == BK * D);
|
||||
static_assert(BQ / 16 < 16); // Barrier 0 joins the CTA; 1..8 join query pairs.
|
||||
constexpr int PREFETCH_VECTORS = (BK * D + THREADS * 8 - 1) / (THREADS * 8);
|
||||
static_assert(THREADS * PREFETCH_VECTORS * 8 >= BK * D);
|
||||
static_assert(KEY_WARPS == 1 && BK == 64);
|
||||
constexpr int shared_bytes() {
|
||||
constexpr int QLD = D + 8, PLD = BK + 4;
|
||||
return (BQ * D + BK * QLD + D * VLD + BQ * PLD) * 2 +
|
||||
(BQ * KEY_WARPS * 2 + BQ * 2) * 4;
|
||||
}
|
||||
|
||||
// QK maxima and probability rows are consumed only by the matching pair.
|
||||
// Keep CTA-wide barriers around K/V staging and tile consumption.
|
||||
__device__ __forceinline__ void sync_query_pair(int query_group) {
|
||||
asm volatile("bar.sync %0, 64;" ::"r"(query_group + 1) : "memory");
|
||||
constexpr int QLD = D + 8;
|
||||
return (BQ * D + BK * QLD + D * VLD) * 2;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int q_swizzle(int row) {
|
||||
@@ -49,19 +41,28 @@ __device__ __forceinline__ void load_q_fragment(fi::AFragment& fragment,
|
||||
values[1] = *reinterpret_cast<const uint4*>(base + ((col + 8) ^ mask));
|
||||
}
|
||||
|
||||
// A 68-half P stride makes accumulator pair stores conflict-free. Odd rows
|
||||
// remain 8-byte aligned, so use 64-bit loads rather than WMMA's 128-bit loads.
|
||||
__device__ __forceinline__ void load_p_fragment(fi::AFragment& fragment,
|
||||
const half* source, int row,
|
||||
int col) {
|
||||
// Convert FP16-rounded probability accumulator pairs to the existing Volta
|
||||
// A fragment layout. Exchange row ownership across lane bit 1, then concatenate
|
||||
// the two eight-column halves across lane bit 3. No arithmetic in the exchange.
|
||||
__device__ __forceinline__ void load_probability_fragment(
|
||||
fi::AFragment& fragment, const unsigned* pairs) {
|
||||
const int lane = threadIdx.x & 31;
|
||||
const int physical_row =
|
||||
row + (lane & 3) + ((lane & 16) >> 2) + ((lane & 4) << 1);
|
||||
const half* base = source + physical_row * (BK + 4) + col;
|
||||
auto* values = reinterpret_cast<uint2*>(fragment.x);
|
||||
const int row_bit = (lane >> 1) & 1;
|
||||
const unsigned own0 = row_bit ? pairs[1] : pairs[0];
|
||||
const unsigned own1 = row_bit ? pairs[3] : pairs[2];
|
||||
const unsigned other0 =
|
||||
__shfl_xor_sync(0xffffffff, row_bit ? pairs[0] : pairs[1], 2);
|
||||
const unsigned other1 =
|
||||
__shfl_xor_sync(0xffffffff, row_bit ? pairs[2] : pairs[3], 2);
|
||||
const unsigned local[4] = {row_bit ? other0 : own0, row_bit ? own0 : other0,
|
||||
row_bit ? other1 : own1, row_bit ? own1 : other1};
|
||||
auto* output = reinterpret_cast<unsigned*>(fragment.x);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
values[i] = *reinterpret_cast<const uint2*>(base + i * 4);
|
||||
for (int part = 0; part < 4; ++part) {
|
||||
const unsigned opposite = __shfl_xor_sync(0xffffffff, local[part], 8);
|
||||
output[part] = (lane & 8) ? opposite : local[part];
|
||||
output[part + 4] = (lane & 8) ? local[part] : opposite;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(THREADS,
|
||||
@@ -69,15 +70,13 @@ __global__ __launch_bounds__(THREADS,
|
||||
const half* v, half* output,
|
||||
int length, int heads,
|
||||
float scale) {
|
||||
constexpr int QLD = D + 8, PLD = BK + 4;
|
||||
constexpr int QLD = D + 8;
|
||||
extern __shared__ __align__(32) unsigned char raw[];
|
||||
half* qs = reinterpret_cast<half*>(raw);
|
||||
half* ks = qs + BQ * D;
|
||||
half* vs = ks + BK * QLD;
|
||||
half* probabilities = vs + D * VLD;
|
||||
float* scores = reinterpret_cast<float*>(probabilities + BQ * PLD);
|
||||
float* maximum = scores + BQ * KEY_WARPS * 2;
|
||||
float* denominator = maximum + BQ;
|
||||
float running_max[2] = {-INFINITY, -INFINITY};
|
||||
float running_sum[2] = {0.f, 0.f};
|
||||
const int tid = threadIdx.x, warp = tid / 32;
|
||||
const int warp_q = warp / KEY_WARPS, warp_k = warp % KEY_WARPS;
|
||||
const int lane = tid % 32;
|
||||
@@ -101,10 +100,6 @@ __global__ __launch_bounds__(THREADS,
|
||||
row < length ? q[base + int64_t(row) * heads * D + i % D]
|
||||
: __float2half(0.f);
|
||||
}
|
||||
if (tid < BQ) {
|
||||
maximum[tid] = -INFINITY;
|
||||
denominator[tid] = 0.f;
|
||||
}
|
||||
__syncthreads();
|
||||
for (int start = 0; start < length; start += BK) {
|
||||
if (start == 0) {
|
||||
@@ -135,76 +130,73 @@ __global__ __launch_bounds__(THREADS,
|
||||
fi::mma_sync_m16n16k16_row_col_f16f16f32(qk[n], qa, kb);
|
||||
}
|
||||
}
|
||||
unsigned probability_pairs[KEY_FRAGMENTS][4];
|
||||
{
|
||||
// Volta distributes each accumulator row across lanes differing in
|
||||
// bits 1 and 3. Reduce its 16 columns in registers, then combine only
|
||||
// the two warp partials through shared memory.
|
||||
float row_max[2] = {-INFINITY, -INFINITY};
|
||||
// Preserve the original two logical K32 partials and their FP32 sum
|
||||
// order even though one warp now owns both halves of this K64 tile.
|
||||
float row_max[2][2] = {{-INFINITY, -INFINITY}, {-INFINITY, -INFINITY}};
|
||||
#pragma unroll
|
||||
for (int n = 0; n < KEY_FRAGMENTS; ++n) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < qk[n].num_elements; ++i) {
|
||||
const int col = warp_k * (BK / KEY_WARPS) + n * 16 + fragment_col +
|
||||
(i & 1) + ((i >> 2) & 1) * 4;
|
||||
const int col = n * 16 + fragment_col + (i & 1) + ((i >> 2) & 1) * 4;
|
||||
const int row = (i >> 1) & 1;
|
||||
qk[n].x[i] = start + col < length ? qk[n].x[i] * scale : -INFINITY;
|
||||
row_max[(i >> 1) & 1] = fmaxf(row_max[(i >> 1) & 1], qk[n].x[i]);
|
||||
row_max[n / 2][row] = fmaxf(row_max[n / 2][row], qk[n].x[i]);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 2; ++r) {
|
||||
row_max[r] =
|
||||
fmaxf(row_max[r], __shfl_xor_sync(0xffffffff, row_max[r], 2));
|
||||
row_max[r] =
|
||||
fmaxf(row_max[r], __shfl_xor_sync(0xffffffff, row_max[r], 8));
|
||||
const int row = warp_q * 16 + fragment_row + r * 2;
|
||||
if ((lane & 10) == 0) scores[row * KEY_WARPS + warp_k] = row_max[r];
|
||||
for (int half = 0; half < 2; ++half) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
auto& maximum = row_max[half][row];
|
||||
maximum = fmaxf(maximum, __shfl_xor_sync(0xffffffff, maximum, 2));
|
||||
maximum = fmaxf(maximum, __shfl_xor_sync(0xffffffff, maximum, 8));
|
||||
}
|
||||
}
|
||||
sync_query_pair(warp_q);
|
||||
float new_max[2], row_sum[2] = {0.f, 0.f};
|
||||
float new_max[2], row_sum[2][2] = {{0.f, 0.f}, {0.f, 0.f}};
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 2; ++r) {
|
||||
const int row = warp_q * 16 + fragment_row + r * 2;
|
||||
new_max[r] = maximum[row];
|
||||
#pragma unroll
|
||||
for (int w = 0; w < KEY_WARPS; ++w)
|
||||
new_max[r] = fmaxf(new_max[r], scores[row * KEY_WARPS + w]);
|
||||
register_alpha[r] = __expf(maximum[row] - new_max[r]);
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
new_max[row] =
|
||||
fmaxf(fmaxf(running_max[row], row_max[0][row]), row_max[1][row]);
|
||||
register_alpha[row] = __expf(running_max[row] - new_max[row]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int n = 0; n < KEY_FRAGMENTS; ++n) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < qk[n].num_elements; ++i) {
|
||||
const int r = (i >> 1) & 1;
|
||||
const int row = warp_q * 16 + fragment_row + r * 2;
|
||||
const int col = warp_k * (BK / KEY_WARPS) + n * 16 + fragment_col +
|
||||
(i & 1) + ((i >> 2) & 1) * 4;
|
||||
const float p = __expf(qk[n].x[i] - new_max[r]);
|
||||
probabilities[row * PLD + col] = __float2half_rn(p);
|
||||
row_sum[r] += p;
|
||||
const int row = (i >> 1) & 1;
|
||||
const float p = __expf(qk[n].x[i] - new_max[row]);
|
||||
qk[n].x[i] = p;
|
||||
row_sum[n / 2][row] += p;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
union PackedPair {
|
||||
half2 value;
|
||||
unsigned bits;
|
||||
} pair;
|
||||
pair.value = __floats2half2_rn(qk[n].x[2 * i], qk[n].x[2 * i + 1]);
|
||||
probability_pairs[n][i] = pair.bits;
|
||||
}
|
||||
}
|
||||
float* partial_sums = scores + BQ * KEY_WARPS;
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 2; ++r) {
|
||||
row_sum[r] += __shfl_xor_sync(0xffffffff, row_sum[r], 2);
|
||||
row_sum[r] += __shfl_xor_sync(0xffffffff, row_sum[r], 8);
|
||||
const int row = warp_q * 16 + fragment_row + r * 2;
|
||||
if ((lane & 10) == 0)
|
||||
partial_sums[row * KEY_WARPS + warp_k] = row_sum[r];
|
||||
}
|
||||
sync_query_pair(warp_q);
|
||||
if (warp_k == 0 && (lane & 10) == 0) {
|
||||
for (int half = 0; half < 2; ++half) {
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 2; ++r) {
|
||||
const int row = warp_q * 16 + fragment_row + r * 2;
|
||||
float sum = 0.f;
|
||||
#pragma unroll
|
||||
for (int w = 0; w < KEY_WARPS; ++w)
|
||||
sum += partial_sums[row * KEY_WARPS + w];
|
||||
denominator[row] = denominator[row] * register_alpha[r] + sum;
|
||||
maximum[row] = new_max[r];
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
auto& sum = row_sum[half][row];
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 2);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 8);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
float sum = 0.f;
|
||||
sum += row_sum[0][row];
|
||||
sum += row_sum[1][row];
|
||||
running_sum[row] = running_sum[row] * register_alpha[row] + sum;
|
||||
running_max[row] = new_max[row];
|
||||
}
|
||||
}
|
||||
// Issue the next K/V global loads while the current V tile is consumed.
|
||||
union StagedVector {
|
||||
@@ -219,7 +211,7 @@ __global__ __launch_bounds__(THREADS,
|
||||
const int tile_row = (tile / (D / 32)) * 8 + (lane >> 2);
|
||||
const int next_row = start + BK + tile_row;
|
||||
const int next_col = (tile % (D / 32)) * 32 + (lane & 3) * 8;
|
||||
if (next_row < length) {
|
||||
if (tile_row < BK && next_row < length) {
|
||||
const int64_t position =
|
||||
base + int64_t(next_row) * heads * D + next_col;
|
||||
if ((reinterpret_cast<uintptr_t>(k) % 16 == 0) &&
|
||||
@@ -247,19 +239,19 @@ __global__ __launch_bounds__(THREADS,
|
||||
}
|
||||
#pragma unroll
|
||||
for (int part = 0; part < OUTPUT_FRAGMENTS; ++part) {
|
||||
const int col = warp_k * (D / KEY_WARPS) + part * 16;
|
||||
const int row = warp_q * 16;
|
||||
auto& pv = accumulators[part];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < pv.num_elements; ++i)
|
||||
pv.x[i] *= register_alpha[(i >> 1) & 1];
|
||||
for (int i = 0; i < accumulators[part].num_elements; ++i)
|
||||
accumulators[part].x[i] *= register_alpha[(i >> 1) & 1];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kv = 0; kv < BK; kv += 16) {
|
||||
fi::AFragment pa;
|
||||
for (int kv = 0; kv < KEY_FRAGMENTS; ++kv) {
|
||||
fi::AFragment pa;
|
||||
load_probability_fragment(pa, probability_pairs[kv]);
|
||||
#pragma unroll
|
||||
for (int part = 0; part < OUTPUT_FRAGMENTS; ++part) {
|
||||
fi::QKBFragment vb;
|
||||
load_p_fragment(pa, probabilities, row, kv);
|
||||
fi::load_qk_b_fragment(vb, vs + col * VLD + kv, VLD);
|
||||
fi::mma_sync_m16n16k16_row_col_f16f16f32(pv, pa, vb);
|
||||
fi::load_qk_b_fragment(vb, vs + part * 16 * VLD + kv * 16, VLD);
|
||||
fi::mma_sync_m16n16k16_row_col_f16f16f32(accumulators[part], pa, vb);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
@@ -269,6 +261,7 @@ __global__ __launch_bounds__(THREADS,
|
||||
const int tile = (tid + n * THREADS) / 32;
|
||||
const int tile_row = (tile / (D / 32)) * 8 + (lane >> 2);
|
||||
const int next_col = (tile % (D / 32)) * 32 + (lane & 3) * 8;
|
||||
if (tile_row >= BK) continue;
|
||||
*reinterpret_cast<uint4*>(ks + tile_row * QLD + next_col) =
|
||||
next_k[n].packed;
|
||||
// Transpose four rows with exact 32-bit lane exchanges. Each lane
|
||||
@@ -292,7 +285,8 @@ __global__ __launch_bounds__(THREADS,
|
||||
__shfl_xor_sync(0xffffffff, transposed[2 * j], 8);
|
||||
const unsigned other1 =
|
||||
__shfl_xor_sync(0xffffffff, transposed[2 * j + 1], 8);
|
||||
const unsigned local = transposed[2 * j + ((lane & 8) >> 3)];
|
||||
const unsigned local =
|
||||
((lane & 8) ? transposed[2 * j + 1] : transposed[2 * j]);
|
||||
const unsigned other = (lane & 8) ? other1 : other0;
|
||||
const uint2 vector =
|
||||
(lane & 8) ? make_uint2(other, local) : make_uint2(local, other);
|
||||
@@ -313,7 +307,7 @@ __global__ __launch_bounds__(THREADS,
|
||||
(i & 1) + ((i >> 2) & 1) * 4;
|
||||
if (q_start + row < length)
|
||||
output[base + int64_t(q_start + row) * heads * D + col] =
|
||||
__float2half_rn(pv.x[i] / denominator[row]);
|
||||
__float2half_rn(pv.x[i] / running_sum[(i >> 1) & 1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -178,7 +178,7 @@ def test_attention_padding_excludes_poisoned_suffix(used, padded):
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU")
|
||||
@pytest.mark.parametrize("length", [127, 128, 129, 12323])
|
||||
@pytest.mark.parametrize("length", [127, 128, 129, 191, 192, 193, 12323])
|
||||
def test_flashinfer_online_softmax_across_tiles_and_batches(length):
|
||||
from vllm.model_executor.models.minimax_h3.cuda_ops import flashinfer_extension
|
||||
|
||||
@@ -197,7 +197,9 @@ def test_flashinfer_online_softmax_across_tiles_and_batches(length):
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU")
|
||||
@pytest.mark.parametrize("length", [31, 32, 33, 63, 64, 65, 127, 128, 129])
|
||||
@pytest.mark.parametrize(
|
||||
"length", [31, 32, 33, 63, 64, 65, 127, 128, 129, 191, 192, 193, 385]
|
||||
)
|
||||
def test_flashinfer_prefetch_tail_and_unaligned_storage(length):
|
||||
from vllm.model_executor.models.minimax_h3.cuda_ops import flashinfer_extension
|
||||
|
||||
@@ -218,7 +220,7 @@ def test_flashinfer_prefetch_tail_and_unaligned_storage(length):
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU")
|
||||
@pytest.mark.parametrize("length", [129, 257])
|
||||
@pytest.mark.parametrize("length", [129, 193, 257, 385])
|
||||
def test_flashinfer_query_groups_have_independent_softmax_state(length):
|
||||
from vllm.model_executor.models.minimax_h3.cuda_ops import flashinfer_extension
|
||||
|
||||
|
||||
Reference in new issue
Block a user