Export torch_scaled_dot_product_attention

[Copilot]

torch_scaled_dot_product_attention was implemented but not exported, forcing users to access it via get(..., envir = asNamespace("torch")). This prevented straightforward use of the fused CUDA kernels that provide 2-3x speedup over manual attention computation.

Changes

• R/gen-namespace-docs.R: Added roxygen documentation with @export tag covering all 8 parameters, mathematical formula, and examples for basic usage, causal masking, and custom attention masks
• NAMESPACE: Added export declaration
• tests/testthat/test-gen-namespace.R: Added tests for basic usage, causal masking, attention masks, and dropout

Usage

library(torch)

query <- torch_randn(2, 8, 10, 64)  # (batch, heads, seq_len, dim)
key <- torch_randn(2, 8, 10, 64)
value <- torch_randn(2, 8, 10, 64)

# Now accessible directly
output <- torch_scaled_dot_product_attention(query, key, value)

# Supports causal masking for autoregressive models
output <- torch_scaled_dot_product_attention(query, key, value, is_causal = TRUE)

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Export torch_scaled_dot_product_attention (SDPA)</issue_title>
<issue_description>## Summary

torch_scaled_dot_product_attention exists in the torch namespace but is not exported. Exporting it would provide significant performance benefits for transformer inference.

Current Workaround

# Have to access unexported function
sdpa <- get("torch_scaled_dot_product_attention", envir = asNamespace("torch"))
output <- sdpa(query, key, value, attn_mask = mask, is_causal = FALSE)

Why This Matters

SDPA uses fused CUDA kernels that are 2-3x faster than manual attention:

# Manual attention
scores <- torch_matmul(q, k$transpose(-2, -1)) / sqrt(head_dim)
attn_weights <- nnf_softmax(scores, dim = -1)
output <- torch_matmul(attn_weights, v)

# SDPA (fused kernel, 2.7x faster in benchmarks)
output <- torch_scaled_dot_product_attention(q, k, v)

For transformer models with 30+ layers, this adds up to meaningful speedups.

Proposed Change

Export torch_scaled_dot_product_attention in the package NAMESPACE so users can call it directly:

torch::torch_scaled_dot_product_attention(query, key, value)

Environment

• torch version: 0.13.0
• R version: 4.4.x
• Discovered while porting Chatterbox TTS to native R torch</issue_description>

<agent_instructions>To export, make sure you add some docs and the #' @export tag. Also make sure the docs are not in the file that's autogenerated. Also please add a test !</agent_instructions>

Comments on the Issue (you are @copilot in this section)

• Fixes Export torch_scaled_dot_product_attention (SDPA) #1403

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

[dfalbel]

lg