Skip to content

[WIP] [fusilli] Emit flex_attention for SDPA #413

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
keshavvinayak01 wants to merge 6 commits into
base: main
Choose a base branch
from
Draft

[WIP] [fusilli] Emit flex_attention for SDPA #413

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
4066b13
[fusilli] Emit flex_attention for SDPA
keshavvinayak01 May 12, 2026
d94ec7e
Merge branch 'main' into issue-275-generate-stats
keshavvinayak01 May 13, 2026
64ddd2c
[fusilli] Fix SDPA flex attention CI failures
keshavvinayak01 May 13, 2026
0bb3c82
[fusilli] Emit causal SDPA through flex attention
keshavvinayak01 May 13, 2026
743a400
[fusilli] Use flex attention for custom scale SDPA custom op
keshavvinayak01 May 13, 2026
5a49893
[fusilli] Emit custom scale SDPA through flex attention
keshavvinayak01 May 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion include/fusilli/attributes/sdpa_attributes.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ class SdpaAttr : public AttributesCRTP<SdpaAttr> {
public:
// Names for Tensor Inputs and Outputs.
enum class InputNames : uint8_t { Q, K, V, MASK };
enum class OutputNames : uint8_t { O };
enum class OutputNames : uint8_t { O, STATS };

std::unordered_map<InputNames, std::shared_ptr<TensorAttr>> inputs;
std::unordered_map<OutputNames, std::shared_ptr<TensorAttr>> outputs;
Expand All @@ -39,13 +39,15 @@ class SdpaAttr : public AttributesCRTP<SdpaAttr> {
FUSILLI_GENERIC_INPUT_TENSOR_SETTER(SdpaAttr, InputNames, V)
FUSILLI_GENERIC_INPUT_TENSOR_SETTER(SdpaAttr, InputNames, MASK)
FUSILLI_GENERIC_OUTPUT_TENSOR_SETTER(SdpaAttr, OutputNames, O)
FUSILLI_GENERIC_OUTPUT_TENSOR_SETTER(SdpaAttr, OutputNames, STATS)

// Tensor getters:
FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, Q)
FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, K)
FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, V)
FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, MASK)
FUSILLI_GENERIC_OUTPUT_TENSOR_GETTER(OutputNames, O)
FUSILLI_GENERIC_OUTPUT_TENSOR_GETTER(OutputNames, STATS)

// Scalar attribute setters:
SdpaAttr &setDropout(float p) {
Expand All @@ -68,17 +70,24 @@ class SdpaAttr : public AttributesCRTP<SdpaAttr> {
return *this;
}

SdpaAttr &setGenerateStats(bool v) {
generateStats_ = v;
return *this;
}

// Scalar attribute getters:
float getDropout() const { return dropout_; }
bool getIsCausal() const { return isCausal_; }
std::optional<float> getScale() const { return scale_; }
bool getEnableGqa() const { return enableGqa_; }
bool getGenerateStats() const { return generateStats_; }

private:
float dropout_ = 0.0f;
bool isCausal_ = false;
std::optional<float> scale_ = std::nullopt;
bool enableGqa_ = false;
bool generateStats_ = false;
};

} // namespace fusilli
Expand Down
11 changes: 9 additions & 2 deletions include/fusilli/graph/graph.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1186,10 +1186,17 @@ Graph::sdpa(const std::shared_ptr<TensorAttr> &q,
// Set outputs.
auto o = outputTensor(sdpaAttr.getName() + "_O");
sdpaAttr.setO(o);
if (sdpaAttr.getGenerateStats() && !sdpaAttr.getSTATS()) {
auto stats = outputTensor(sdpaAttr.getName() + "_STATS");
stats->setDataType(DataType::Float);
sdpaAttr.setSTATS(stats);
}

// Create node and add to Graph's subNodes_.
// Keep sdpaAttr populated for callers: Graph::sdpa returns O, and the
// generate_stats STATS tensor is exposed through sdpaAttr.
SdpaAttr nodeAttr = sdpaAttr;
subNodes_.emplace_back(
std::make_unique<SdpaNode>(std::move(sdpaAttr), context));
std::make_unique<SdpaNode>(std::move(nodeAttr), context));

return o;
}
Expand Down
45 changes: 45 additions & 0 deletions include/fusilli/node/sdpa_node.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,15 +41,24 @@ class SdpaNode : public NodeCRTP<SdpaNode> {
: NodeCRTP(ctx), sdpaAttr(std::move(attr)) {}

// ASM emitter methods.
std::string emitModuleScopeAsm() const override final;
std::string emitNodePreAsm() const override final;
std::string getOperandNamesAsm() const;
std::string getOperandTypesAsm() const;
std::string getFlexAttnOperandNamesAsm() const;
std::string getFlexAttnOperandTypesAsm() const;
std::string getResultNamesAsm() const;
std::string getResultTypesAsm() const;
std::string getFlexAttnResultNamesAsm() const;
std::string getFlexAttnResultTypesAsm() const;
std::string getDropoutOpsAsm() const;
std::string getIsCausalOpsAsm() const;
std::string getScaleOpsAsm() const;
std::string getEnableGqaOpsAsm() const;
std::string getReturnLseOpsAsm() const;
std::string getReturnMaxScoresOpsAsm() const;
std::string getCausalMaskFnNameAsm() const;
bool useLegacySdpaAsm() const;

const std::string &getName() const override final {
return sdpaAttr.getName();
Expand All @@ -64,6 +73,7 @@ class SdpaNode : public NodeCRTP<SdpaNode> {
std::shared_ptr<TensorAttr> kT = sdpaAttr.getK();
std::shared_ptr<TensorAttr> vT = sdpaAttr.getV();
std::shared_ptr<TensorAttr> oT = sdpaAttr.getO();
std::shared_ptr<TensorAttr> statsT = sdpaAttr.getSTATS();
std::shared_ptr<TensorAttr> maskT = sdpaAttr.getMASK();

// Ensure mandatory input and output tensors are set.
Expand All @@ -75,6 +85,16 @@ class SdpaNode : public NodeCRTP<SdpaNode> {
"SDPA input tensor V not set");
FUSILLI_RETURN_ERROR_IF(!oT, ErrorCode::AttributeNotSet,
"SDPA output tensor O not set");
if (sdpaAttr.getGenerateStats()) {
FUSILLI_RETURN_ERROR_IF(
!statsT, ErrorCode::AttributeNotSet,
"SDPA output tensor STATS not set when generate_stats is enabled");
} else {
FUSILLI_RETURN_ERROR_IF(
statsT, ErrorCode::InvalidAttribute,
"SDPA output tensor STATS should not be set when generate_stats is "
"disabled");
}

// Rank checks: all tensors must be rank 4.
constexpr size_t kRequiredRank = 4;
Expand Down Expand Up @@ -192,6 +212,7 @@ class SdpaNode : public NodeCRTP<SdpaNode> {
std::shared_ptr<TensorAttr> qT = sdpaAttr.getQ();
std::shared_ptr<TensorAttr> vT = sdpaAttr.getV();
std::shared_ptr<TensorAttr> oT = sdpaAttr.getO();
std::shared_ptr<TensorAttr> statsT = sdpaAttr.getSTATS();

const std::vector<int64_t> &qDim = qT->getDim();
const std::vector<int64_t> &vDim = vT->getDim();
Expand All @@ -209,6 +230,18 @@ class SdpaNode : public NodeCRTP<SdpaNode> {
generateStrideFromDim(oDim, getContiguousStrideOrder(oDim.size())));
}

if (statsT) {
std::vector<int64_t> statsDim = {qDim[0], qDim[1], qDim[2]};
if (statsT->getDim().empty())
statsT->setDim(statsDim);
if (statsT->getStride().empty()) {
statsT->setStride(generateStrideFromDim(
statsDim, getContiguousStrideOrder(statsDim.size())));
}
if (statsT->getDataType() == DataType::NotSet)
statsT->setDataType(DataType::Float);
}

return ok();
}

Expand All @@ -219,6 +252,7 @@ class SdpaNode : public NodeCRTP<SdpaNode> {
std::shared_ptr<TensorAttr> qT = sdpaAttr.getQ();
std::shared_ptr<TensorAttr> vT = sdpaAttr.getV();
std::shared_ptr<TensorAttr> oT = sdpaAttr.getO();
std::shared_ptr<TensorAttr> statsT = sdpaAttr.getSTATS();

const std::vector<int64_t> &qDim = qT->getDim();
const std::vector<int64_t> &vDim = vT->getDim();
Expand All @@ -231,6 +265,17 @@ class SdpaNode : public NodeCRTP<SdpaNode> {
"SDPA output tensor O dimensions do not match expected shape "
"[batch, headsQ, seqQ, headDim]");

if (sdpaAttr.getGenerateStats()) {
std::vector<int64_t> expectedStatsDim = {qDim[0], qDim[1], qDim[2]};
FUSILLI_RETURN_ERROR_IF(
statsT->getDim() != expectedStatsDim, ErrorCode::InvalidAttribute,
"SDPA output tensor STATS dimensions do not match expected shape "
"[batch, headsQ, seqQ]");
FUSILLI_RETURN_ERROR_IF(
statsT->getDataType() != DataType::Float, ErrorCode::InvalidAttribute,
"SDPA output tensor STATS must have Float data type");
}

return ok();
}
};
Expand Down
Loading
Loading