…lines 1–256 unchanged (256 lines)
257 window_size_right: int | None,
258 seqused_k: torch.Tensor | None,
259 out: torch.Tensor | None = None,
260 q_descale: torch.Tensor | None = None,
261 k_descale: torch.Tensor | None = None,
262 v_descale: torch.Tensor | None = None,
263 block_table: torch.Tensor | None = None,
264) -> tuple[torch.Tensor, torch.Tensor]:
265 """
266 Run the FA3 forward pass by calling the C++ kernel directly.
267 """
268 if _FA3_CUDA_FWD is None:
269 raise RuntimeError("FA3 not registered")
…lines 270–277 unchanged (8 lines)
278 else _maybe_contiguous(value)
279 )
280
281 cu_seqlens_q = _maybe_contiguous(cu_seq_q)
282 cu_seqlens_k = _maybe_contiguous(cu_seq_k)
283 seqused_k = _maybe_contiguous(seqused_k)
284 block_table = _maybe_contiguous(block_table)
285
286 out, softmax_lse, out_accum, softmax_lse_accum = _FA3_CUDA_FWD(
287 q,
288 k,
289 v,
290 None, # k_new
…lines 291–294 unchanged (4 lines)
295 cu_seqlens_k, # cu_seqlens_k
296 None, # cu_seqlens_k_new
297 None, # seqused_q
298 seqused_k, # seqused_k
299 max_q, # max_seqlen_q
300 max_k, # max_seqlen_k
301 block_table, # block_table,
302 None, # kv_batch_idx,
303 None, # leftpad_k,
304 None, # rotary_cos,
305 None, # rotary_sin,
306 None, # seqlens_rotary,
307 q_descale, # q_descale,
…lines 308–398 unchanged (91 lines)
399 scale: float | None = None,
400 window_size_left: int = -1,
401 window_size_right: int = -1,
402 seqused_k: torch.Tensor | None = None,
403 alibi_slopes: torch.Tensor | None = None,
404 out: torch.Tensor | None = None,
405 block_table: torch.Tensor | None = None,
406):
407 error = _fa3_forward_support_error(
408 query,
409 key,
410 value,
411 dropout_p,
…lines 412–432 unchanged (21 lines)
433 window_size_right,
434 seqused_k,
435 out,
436 q_descale,
437 k_descale,
438 v_descale,
439 block_table,
440 )
441 rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
442 philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
443 debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
444 return out, lse, rng_state, philox_offset, debug_mask
445
…lines 446–457 unchanged (12 lines)
458 *,
459 scale: float | None = None,
460 window_size_left: int = -1,
461 window_size_right: int = -1,
462 seqused_k: torch.Tensor | None = None,
463 alibi_slopes: torch.Tensor | None = None,
464 block_table: torch.Tensor | None = None,
465 out: torch.Tensor | None = None,
466):
467 return _fa3_flash_attention_forward_impl(
468 query,
469 key,
470 value,
…lines 471–480 unchanged (10 lines)
481 scale=scale,
482 window_size_left=window_size_left,
483 window_size_right=window_size_right,
484 seqused_k=seqused_k,
485 alibi_slopes=alibi_slopes,
486 out=out,
487 block_table=block_table,
488 )
489
490
491def _fa3_flash_attention_backward_impl(
492 grad_out: torch.Tensor,
493 query: torch.Tensor,
…lines 494–699 unchanged (206 lines)